<commit_before>// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <algorithm>
#include <numeric>

#include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOpUtils.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
#include "iree/compiler/Dialect/IREE/IR/IREETypes.h"
#include "iree/compiler/Dialect/Shape/IR/ShapeOps.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/Optional.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/StringExtras.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/IR/Value.h"
#include "mlir/Support/LogicalResult.h"

namespace mlir {
namespace iree_compiler {
namespace IREE {
namespace Flow {

//===----------------------------------------------------------------------===//
// Folding utilities
//===----------------------------------------------------------------------===//

// Returns a new set of dynamic dimensions for a shape carrying op when a type
// is being changed. This attempts to reuse the existing dimension values if
// they are available and will drop/insert new ones as required.
static SmallVector<Value, 4> refreshDimsOnTypeChange(
    Operation *op, Type oldType, Type newType, ValueRange oldDims,
    PatternRewriter &rewriter) {
  if (oldType == newType) return llvm::to_vector<4>(oldDims);

  // Build an expanded list of all the dims - constants will be nullptr.
  // This lets us map back the new types without worrying about whether some
  // subset become static or dynamic.
  auto oldShapedType = oldType.cast<ShapedType>();
  SmallVector<Value, 4> allOldDims(oldShapedType.getRank());
  for (unsigned i = 0; i < oldShapedType.getRank(); ++i) {
    if (oldShapedType.isDynamicDim(i)) {
      allOldDims[i] = oldDims.front();
      oldDims = oldDims.drop_front();
    }
  }

  auto newShapedType = newType.cast<ShapedType>();
  SmallVector<Value, 4> newDims;
  for (unsigned i = 0; i < newShapedType.getRank(); ++i) {
    if (newShapedType.isDynamicDim(i)) {
      auto oldValue = allOldDims[i];
      if (oldValue) {
        // Old value valid; reuse.
        newDims.push_back(oldValue);
      } else {
        // Dimension has changed to be dynamic; insert a constant to use.
        // This sometimes happens during folding of casts and usually is cleaned
        // up pretty quickly.
        newDims.push_back(rewriter.createOrFold<ConstantIndexOp>(
            op->getLoc(), oldShapedType.getDimSize(i)));
      }
    }
  }
  return newDims;
}

//===----------------------------------------------------------------------===//
// Streams
//===----------------------------------------------------------------------===//

namespace {

// Returns true if the given |value| is used again after |updateOp| consumes it.
static bool hasUsersInStreamAfterUpdate(Value value, Operation *updateOp) {
  for (auto user : value.getUsers()) {
    if (user == updateOp) continue;
    if (user->getBlock() != updateOp->getBlock() ||
        user->isBeforeInBlock(updateOp)) {
      // From a dominating block or earlier in the block, cannot be a consumer.
      continue;
    }
    return true;
  }
  return false;
}

/// Inserts clones into the stream as required by tied results.
/// This is required to preserve the immutable tensor semantics required by the
/// SSA use-def chain.
///
/// Example:
///   %0 = flow.dispatch
///   // %0 will be updated in-place and renamed %1:
///   %1 = flow.dispatch %0 -> %0
///   // The original value of %0 (aka %1) is required but is not valid!
///   %2 = flow.dispatch %0
/// ->
///   %0 = flow.dispatch
///   // Capture the value of %0 before it is modified:
///   %clone = flow.tensor.clone %0
///   // Update %0 in-place and rename to %1, safe as %0 now has one use:
///   %1 = flow.dispatch %0 -> %0
///   // Use the cloned %0 value:
///   %2 = flow.dispatch %clone
struct InsertImmutabilityPreservingStreamClones
    : public OpRewritePattern<ExStreamFragmentOp> {
  using OpRewritePattern<ExStreamFragmentOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(ExStreamFragmentOp op,
                                PatternRewriter &rewriter) const override {
    bool didClone =
        insertTiedClones(cast<TiedOpInterface>(op.getOperation()), rewriter);
    for (auto &block : op.getClosureBodyRegion()) {
      for (auto &innerOp : block) {
        if (auto tiedOp = dyn_cast<TiedOpInterface>(innerOp)) {
          didClone |= insertTiedClones(tiedOp, rewriter);
        }
      }
    }
    return didClone ? success() : failure();
  }

  bool insertTiedClones(TiedOpInterface tiedOp,
                        PatternRewriter &rewriter) const {
    bool didClone = false;
    for (unsigned resultIndex = 0; resultIndex < tiedOp->getNumResults();
         ++resultIndex) {
      auto tiedOperandIndex = tiedOp.getTiedResultOperandIndex(resultIndex);
      if (!tiedOperandIndex.hasValue()) continue;
      auto tiedOperand = tiedOp->getOperand(tiedOperandIndex.getValue());
      if (hasUsersInStreamAfterUpdate(tiedOperand, tiedOp)) {
        rewriter.setInsertionPointAfterValue(tiedOperand);
        auto clonedOperand = rewriter.createOrFold<TensorCloneOp>(
            tiedOperand.getLoc(), tiedOperand);
        SmallPtrSet<Operation *, 1> excludedOps;
        excludedOps.insert(tiedOp.getOperation());
        excludedOps.insert(clonedOperand.getDefiningOp());
        tiedOperand.replaceAllUsesExcept(clonedOperand, excludedOps);
        didClone = true;
      }
    }
    return didClone;
  }
};
<commit_msg>Insert clones prior to mutation and not where it originates. (#5292)<commit_after>// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <algorithm>
#include <numeric>

#include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOpUtils.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
#include "iree/compiler/Dialect/IREE/IR/IREETypes.h"
#include "iree/compiler/Dialect/Shape/IR/ShapeOps.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/Optional.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/StringExtras.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/IR/Value.h"
#include "mlir/Support/LogicalResult.h"

namespace mlir {
namespace iree_compiler {
namespace IREE {
namespace Flow {

//===----------------------------------------------------------------------===//
// Folding utilities
//===----------------------------------------------------------------------===//

// Returns a new set of dynamic dimensions for a shape carrying op when a type
// is being changed. This attempts to reuse the existing dimension values if
// they are available and will drop/insert new ones as required.
static SmallVector<Value, 4> refreshDimsOnTypeChange(
    Operation *op, Type oldType, Type newType, ValueRange oldDims,
    PatternRewriter &rewriter) {
  if (oldType == newType) return llvm::to_vector<4>(oldDims);

  // Build an expanded list of all the dims - constants will be nullptr.
  // This lets us map back the new types without worrying about whether some
  // subset become static or dynamic.
  auto oldShapedType = oldType.cast<ShapedType>();
  SmallVector<Value, 4> allOldDims(oldShapedType.getRank());
  for (unsigned i = 0; i < oldShapedType.getRank(); ++i) {
    if (oldShapedType.isDynamicDim(i)) {
      allOldDims[i] = oldDims.front();
      oldDims = oldDims.drop_front();
    }
  }

  auto newShapedType = newType.cast<ShapedType>();
  SmallVector<Value, 4> newDims;
  for (unsigned i = 0; i < newShapedType.getRank(); ++i) {
    if (newShapedType.isDynamicDim(i)) {
      auto oldValue = allOldDims[i];
      if (oldValue) {
        // Old value valid; reuse.
        newDims.push_back(oldValue);
      } else {
        // Dimension has changed to be dynamic; insert a constant to use.
        // This sometimes happens during folding of casts and usually is cleaned
        // up pretty quickly.
        newDims.push_back(rewriter.createOrFold<ConstantIndexOp>(
            op->getLoc(), oldShapedType.getDimSize(i)));
      }
    }
  }
  return newDims;
}

//===----------------------------------------------------------------------===//
// Streams
//===----------------------------------------------------------------------===//

namespace {

// Returns true if the given |value| is used again after |updateOp| consumes it.
static bool hasUsersInStreamAfterUpdate(Value value, Operation *updateOp) {
  for (auto user : value.getUsers()) {
    if (user == updateOp) continue;
    if (user->getBlock() != updateOp->getBlock() ||
        user->isBeforeInBlock(updateOp)) {
      // From a dominating block or earlier in the block, cannot be a consumer.
      continue;
    }
    return true;
  }
  return false;
}

/// Inserts clones into the stream as required by tied results.
/// This is required to preserve the immutable tensor semantics required by the
/// SSA use-def chain.
///
/// Example:
///   %0 = flow.dispatch
///   // %0 will be updated in-place and renamed %1:
///   %1 = flow.dispatch %0 -> %0
///   // The original value of %0 (aka %1) is required but is not valid!
///   %2 = flow.dispatch %0
/// ->
///   %0 = flow.dispatch
///   // Capture the value of %0 before it is modified:
///   %clone = flow.tensor.clone %0
///   // Update %0 in-place and rename to %1, safe as %0 now has one use:
///   %1 = flow.dispatch %0 -> %0
///   // Use the cloned %0 value:
///   %2 = flow.dispatch %clone
struct InsertImmutabilityPreservingStreamClones
    : public OpRewritePattern<ExStreamFragmentOp> {
  using OpRewritePattern<ExStreamFragmentOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(ExStreamFragmentOp op,
                                PatternRewriter &rewriter) const override {
    bool didClone =
        insertTiedClones(cast<TiedOpInterface>(op.getOperation()), rewriter);
    for (auto &block : op.getClosureBodyRegion()) {
      for (auto &innerOp : block) {
        if (auto tiedOp = dyn_cast<TiedOpInterface>(innerOp)) {
          didClone |= insertTiedClones(tiedOp, rewriter);
        }
      }
    }
    return didClone ? success() : failure();
  }

  bool insertTiedClones(TiedOpInterface tiedOp,
                        PatternRewriter &rewriter) const {
    bool didClone = false;
    for (unsigned resultIndex = 0; resultIndex < tiedOp->getNumResults();
         ++resultIndex) {
      auto tiedOperandIndex = tiedOp.getTiedResultOperandIndex(resultIndex);
      if (!tiedOperandIndex.hasValue()) continue;
      auto tiedOperand = tiedOp->getOperand(tiedOperandIndex.getValue());
      if (hasUsersInStreamAfterUpdate(tiedOperand, tiedOp)) {
        rewriter.setInsertionPoint(tiedOp);
        auto clonedOperand = rewriter.createOrFold<TensorCloneOp>(
            tiedOperand.getLoc(), tiedOperand);
        SmallPtrSet<Operation *, 1> excludedOps;
        excludedOps.insert(tiedOp.getOperation());
        excludedOps.insert(clonedOperand.getDefiningOp());
        tiedOperand.replaceAllUsesExcept(clonedOperand, excludedOps);
        didClone = true;
      }
    }
    return didClone;
  }
};

{'fraction_non_alphanumeric': 0.16285936285936287, 'fraction_numerical': 0.006060606060606061, 'mean_word_length': 2.711361014994233, 'pattern_counts': {'":': 0, '<': 39, '<?xml version=': 0, '>': 51, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7581656', 'n_tokens_mistral': 3917, 'n_tokens_neox': 3876, 'n_words': 1076}