[mlir][linalg] Add support for WAW fusion on tensors.

[gysit]

mlir/test/Dialect/Linalg/fusion-tensor-pattern.mlir

173

Should probably be %{{.*}}, %{{.*}}?

[nicolasvasilache]

Address comments.

[mravishankar]

Couldnt get to this before it landed. Leaving my comments here anyway.

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp

179

Does this really need to be added as a dependence? init_tensor is mostly shape so there should not be any dependence between them....

183

Also interesting that this is added as a WAW dependence. In tensor world this is not a WAW dependence really. it is actually a RAW dependence. Take this example,

%1 = linalg.fill(%0, %cst)
%2 = linalg.matmul ins(%lhs1, %rhs1) outs(%1)
%3 = linalg.matmul ins(%lhs2, %rhs2) outs(%1)

In both these examples the %1 value is "read" by the two matmul operations. It is not written. So it is strange to see this as a WAW dependence. It might just be a matter of terminology clash between tensor world and buffer world, cause a WAW dependence in the traditional sense does not have a place in tensor world (you can write to a tensor twice). Its really a RAW dependence, but with the additional information that when bufferized it becomes a WAW dependence? So concern is that in might be conflating issues, and in general hard to follow.

Do we even need to make this a WAW dependence. Its just a dependence. Why not just say RAW dependence?

[nicolasvasilache]

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp

179

isInitTensor means the value is actually used, so yes.
Without a dependence the fill does not fuse into the matmul, I didn't add these just for making the impl more confusing :) .

183

This is a good point, I'll mull it over the week-end.
The thinking is that this dependence into an output tensor has special semantics, see how the scf.for iter_args is updated below.
I'll add an example for which it is not an init_tensor and see what happens.

mlir/include/mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h

SmallVector<LinalgDependenceGraphElem, 2> getDependentOperationsFrom(
                              DependenceType::RAW, DependenceType::WAW}) const;
 
   /// Returns all dependent operations (into and from) given `operation`.
   SmallVector<LinalgDependenceGraphElem, 2>
   getDependentOperations(LinalgOp linalgOp,
                          ArrayRef<DependenceType> depTypes = {
                              DependenceType::RAW, DependenceType::WAW}) const;
 
+  void print(raw_ostream &os) const;
+
+  void dump() const;
+
 private:
   // Keep dependences in both directions, this is not just a performance gain
   // but it also reduces usage errors.
   // Dependence information is stored as a map of:
   //   (source operation -> LinalgDependenceGraphElem)
   DependenceGraph dependencesFromGraphs[DependenceType::NumTypes];
   // Reverse dependence information is stored as a map of:
   //   (destination operation -> LinalgDependenceGraphElem)

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp

LinalgDependenceGraph::getDependencesInto(
     Operation *dst, LinalgDependenceGraph::DependenceType dt) const {
   auto iter = dependencesIntoGraphs[dt].find(dst);
   if (iter == dependencesIntoGraphs[dt].end())
     return llvm::make_range(nullptr, nullptr);
   return llvm::make_range(iter->second.begin(), iter->second.end());
 }
 
 void LinalgDependenceGraph::addDependencesBetween(LinalgOp src, LinalgOp dst) {
+  LLVM_DEBUG(dbgs() << "addDependencesBetween " << *src.getOperation()
+                    << " and " << *dst.getOperation() << "\n");
   if (src.hasTensorSemantics() && dst.hasTensorSemantics()) {
     for (OpOperand &dstOpOperand : dst.getInputOpOperands()) {
       // Check if the operand is defined by the src.
       auto definingOp = dstOpOperand.get().getDefiningOp<LinalgOp>();
       if (definingOp && definingOp == src)
         addDependenceElem(DependenceType::RAW, dstOpOperand.get(),
                           &dstOpOperand);
     }
+    for (OpOperand &dstOpOperand : dst.getOutputOpOperands()) {
+      // Check if the operand is defined by the src.
+      auto definingOp = dstOpOperand.get().getDefiningOp<LinalgOp>();
+      if (definingOp && definingOp == src) {
+        if (dst.isInitTensor(&dstOpOperand)) {
+          addDependenceElem(DependenceType::RAW, dstOpOperand.get(),
+                            &dstOpOperand);
+        }
+        addDependenceElem(DependenceType::WAW, dstOpOperand.get(),
+                          &dstOpOperand);
+      }
+    }
     return;
   }
   assert(src.hasBufferSemantics() && dst.hasBufferSemantics() &&
          "unhandled dependence tracking for mixed buffer/tensor operations");
   for (OpOperand *srcOpOperand : src.getOutputBuffersOpOperands()) { // W
     // RAW graph
     for (OpOperand *dstOpOperand : dst.getInputBuffersOpOperands()) // R
       if (aliases.alias(srcOpOperand->get(), dstOpOperand->get()))  // RAW alias
LinalgDependenceGraph::getDependentOperations(
     LinalgOp linalgOp, ArrayRef<DependenceType> depTypes) const {
   SmallVector<LinalgDependenceGraphElem, 2> dependentOperations =
       getDependentOperationsInto(linalgOp, depTypes);
   SmallVector<LinalgDependenceGraphElem, 2> t =
       getDependentOperationsFrom(linalgOp, depTypes);
   dependentOperations.append(t.begin(), t.end());
   return dependentOperations;
 }
+
+void LinalgDependenceGraph::print(raw_ostream &os) const {
+  for (auto dt : {
+           LinalgDependenceGraph::DependenceType::RAW,
+           LinalgDependenceGraph::DependenceType::WAW,
+       }) {
+    const auto &fromGraph = dependencesFromGraphs[dt];
+    for (const auto &it : fromGraph) {
+      os << "[LinalgDependenceGraph] DT " << dt << " from: " << *it.first
+         << ":\n";
+      for (const auto &dep : it.second) {
+        os << "\tDT " << dt << " " << *dep.getDependentOp() << ":\n";
+      }
+    }
+  }
+}
+
+void LinalgDependenceGraph::dump() const { print(llvm::errs()); }

mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp

 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
 #include <set>
 
 #define DEBUG_TYPE "linalg-fusion"
 
 using namespace mlir;

 }
 
 /// For `consumer` with buffer semantics, find the Linalg operation on buffers
 /// that is the last writer of `consumerOpOperand`. For now the fusable
 /// dependence is returned as an instance of the `dependenceGraph`.
 static Optional<LinalgDependenceGraph::LinalgDependenceGraphElem>
 findFusableProducer(OpOperand &consumerOpOperand,
                     const LinalgDependenceGraph &dependenceGraph) {
+  LLVM_DEBUG(llvm::dbgs() << "findFusableProducer for: "
+                          << consumerOpOperand.get() << " @"
+                          << consumerOpOperand.getOperandNumber() << " in "
+                          << *consumerOpOperand.getOwner() << "\n");
   LinalgOp consumerOp = dyn_cast<LinalgOp>(consumerOpOperand.getOwner());
   if (!consumerOp)
     return {};
 
   // Only consider RAW and WAW atm.
   for (auto depType : {
            LinalgDependenceGraph::DependenceType::RAW,
            LinalgDependenceGraph::DependenceType::WAW,
        }) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "Dependencies into: " << *consumerOp.getOperation() << "\n");
     for (auto dependence : llvm::make_filter_range(
              dependenceGraph.getDependencesInto(consumerOp, depType),
              [&](LinalgDependenceGraph::LinalgDependenceGraphElem elem) {
+               LLVM_DEBUG(llvm::dbgs() << "Inspect dependence btw: "
+                                       << elem.getIndexingValue() << " and "
+                                       << elem.getDependentValue() << "\n");
                Value v = elem.getIndexingValue();
                Optional<unsigned> operandNum =
                    elem.getIndexingOpViewOperandNum();
                return isa<LinalgOp>(elem.getDependentOp()) &&
                       v == consumerOpOperand.get() && operandNum &&
                       operandNum.getValue() ==
                           consumerOpOperand.getOperandNumber();
              })) {
static Optional<TiledLinalgOp> tileRootOperation(
   return tileLinalgOp(builder, op, tileFusedLoopsOptions);
 }
 
 /// Fuse the operations in `fusionCandidates` with `tiledOp`. Latter is expected
 /// to be a tiled operation such that it is valid to fuse all operations in
 /// `fusionCandidates`, i.e. move the operation within the inter-tile loops of
 /// `tiledOp`.
 static SmallVector<LinalgOp, 1>
-fuseOperations(OpBuilder &builder, LinalgOp rootOp, LinalgOp tiledOp,
+fuseOperations(OpBuilder &builder, LinalgOp rootOp, TiledLinalgOp tiledLinalgOp,
                ArrayRef<LinalgOp> fusionCandidates,
                const FusableOpDependencesTy &fusableDependences,
                const std::set<unsigned> &fusedLoops) {
+  LinalgOp tiledOp = tiledLinalgOp.op;
   OpBuilder::InsertionGuard guard(builder);
   builder.setInsertionPoint(tiledOp);
+
   DenseMap<unsigned, Range> fusedLoopsAndRanges;
   for (unsigned loop : fusedLoops) {
     ShapeDimension shapeDim = getShapeDefiningLoopRange(tiledOp, loop, true);
     fusedLoopsAndRanges[loop] = getRangeFromOperandShape(
         builder, tiledOp.getLoc(), shapeDim.shape, shapeDim.dimension);
   }
 
   SmallVector<LinalgOp, 1> fusedOps(fusionCandidates.size());
   DenseMap<Operation *, LinalgOp> origOpToFusedOp;
   origOpToFusedOp[rootOp.getOperation()] = tiledOp;
   for (auto candidate : enumerate(llvm::reverse(fusionCandidates))) {
     LinalgOp origOp = candidate.value();
     LinalgOp fusedOp = fuse(builder, origOp, fusedLoopsAndRanges);
     origOpToFusedOp[origOp.getOperation()] = fusedOp;
     fusedOps[fusionCandidates.size() - candidate.index() - 1] = fusedOp;
+
+    // Prepare the builder for the next insertion point.
+    auto guard =
+        llvm::make_scope_exit([&]() { builder.setInsertionPoint(fusedOp); });
+    if (!origOp.hasTensorSemantics())
+      continue;
+
     // If the producer consumer operations are linalg operations on tensors, the
     // dependence is due to value produced (as a return tensor) by the producer
     // and used in the consumer. The returned value of the fused op needs to be
     // made the operand of the tiled/fused consumer operation. By construction
     // the value returned by the producer is the value used by the consumer.
     for (auto &dependence : fusableDependences.lookup(origOp.getOperation())) {
-      if (origOp.hasTensorSemantics() &&
-          dependence.dependenceType ==
-              LinalgDependenceGraph::DependenceType::RAW) {
-        unsigned resultIndex =
-            dependence.getDependentOpViewResultNum().getValue();
-        LinalgOp consumer = origOpToFusedOp.lookup(dependence.getIndexingOp());
-        if (!consumer)
-          continue;
-        Value replacementValue = fusedOp.getOperation()->getResult(resultIndex);
-        consumer.getOperation()->setOperand(
-            dependence.getIndexingOpViewOperandNum().getValue(),
-            replacementValue);
-      }
-    }
-    builder.setInsertionPoint(fusedOp);
+      if (dependence.dependenceType !=
+          LinalgDependenceGraph::DependenceType::RAW)
+        continue;
+
+
+
+    // At this point, all Linalg uses of the tensors produced by `origOp` have
+    // been replaced. However, there may still be "output tensor"-like uses
+    // coming from WAW dependencies.
+    // All these uses are iter_args of the outermost loop (TODO: add a check).
+    // Such iter_args uses serve 2 purposes:
+    //  1. give a shape to the output
+    //  2. encode destructive updates that may be inplaceable by bufferization.
+    // To keep the second type of information while letting the unfused op die
+    // unused, we need to forward the producer output operand.
+    for (auto &operand :
+         cast<scf::ForOp>(tiledLinalgOp.loops.front()).getIterOpOperands())
+      if (auto opResult = operand.get().dyn_cast<OpResult>())
+        if (opResult.getOwner() == origOp)
+          operand.set(origOp.getOutputTensors()[opResult.getResultNumber()]);
   }
   return fusedOps;
 }
 
 template <typename LoopType>
 static Optional<TiledAndFusedLinalgOps>
 tileAndFuseLinalgOpsImpl(OpBuilder &builder, ArrayRef<LinalgOp> ops,
                          const LinalgDependenceGraph &dependenceGraph,
if (!tilingOptions.interchangeVector.empty()) {
     return llvm::None;
   }
 
   OpBuilder::InsertionGuard guard(builder);
   builder.setInsertionPoint(rootOp);
   ScopedContext scope(builder, rootOp.getLoc());
 
   // Find all the producers.
+  LLVM_DEBUG(llvm::dbgs() << "findAllFusableDependences\n");
   FusableOpDependencesTy fusableDependences =
       findAllFusableDependences(ops, dependenceGraph);
-  if (fusableDependences.empty())
+  if (fusableDependences.empty()) {
+    LLVM_DEBUG(llvm::dbgs() << "no fusable dependencies found\n");
     return llvm::None;
+  }
 
   TiledAndFusedLinalgOps ret;
   // Find the loops that can be tiled and fused.
+  LLVM_DEBUG(llvm::dbgs() << "collectFusableLoops\n");
   ret.fusedLoopDims = collectFusableLoops(ops, fusableDependences);
 
   // If there are no fusable dependences or there are no tile+fusable loops,
   // just return.
   if (ret.fusedLoopDims.empty()) {
+    LLVM_DEBUG(llvm::dbgs() << "no fusable loops found\n");
     return llvm::None;
   }
 
   // Tile the fused loops in the last operation in the list.
   SmallVector<Value, 4> tileSizeVector =
       tilingOptions.tileSizeComputationFunction(builder, rootOp);
   Optional<TiledLinalgOp> tiledRootOp = tileRootOperation(
       builder, rootOp, tileSizeVector, tilingOptions, ret.fusedLoopDims);
   if (!tiledRootOp) {
     rootOp.emitRemark("failed to tile the fused loops");
     return llvm::None;
   }
   ret.op = tiledRootOp->op;
   ret.fusedLoops.assign(tiledRootOp->loops.begin(), tiledRootOp->loops.end());
 
   // Fuse the other operations into the fused inter-tile loops produced above.
-  ret.fusedProducers = fuseOperations(builder, rootOp, ret.op, ops.drop_back(),
-                                      fusableDependences, ret.fusedLoopDims);
+  ret.fusedProducers =
+      fuseOperations(builder, rootOp, *tiledRootOp, ops.drop_back(),
 
   return ret;
 }
 
 Optional<TiledAndFusedLinalgOps>
 mlir::linalg::tileAndFuseLinalgOps(OpBuilder &builder, ArrayRef<LinalgOp> ops,
                                    const LinalgDependenceGraph &dependenceGraph,
                                    const LinalgTilingOptions &tilingOptions) {

mlir/test/Dialect/Linalg/fusion-tensor-pattern.mlir

 //       CHECK:       %[[ST_RESULT:.+]] = linalg.generic
 //  CHECK-SAME:         ins(%[[LHS]] : tensor<?x?xf32>)
 //  CHECK-SAME:         outs(%[[ST_ARG6]] : tensor<?x?xf32>)
 //       CHECK:       %[[UPDATE:.+]] = subtensor_insert %[[ST_RESULT]]
 //  CHECK-SAME:         into %[[ARG6]][%[[IV0]], %[[IV1]]]
 //       CHECK:       scf.yield %[[UPDATE]]
 //       CHECK:     scf.yield %[[YIELD]]
 //       CHECK:   return %[[RESULT]]
+
+// -----
+
+module {
+  func @matmul_out_fusion(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
+                      %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {
+    %c0 = constant 0.0 : f32
+    %0 = linalg.fill(%arg0, %c0) : tensor<?x?xf32>, f32 -> tensor<?x?xf32>
+    %1 = linalg.matmul {__internal_linalg_transform__ = "out_fusion"}
+      ins(%arg1, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32>
+    return %1 : tensor<?x?xf32>
+  }
+}
+
+// CHECK-LABEL: func @matmul_out_fusion(
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//       CHECK: %[[C0:.*]] = constant 0.0{{.*}} : f32
+//   CHECK-NOT: fill
+//       CHECK: scf.for %[[I:.*]]{{.*}}iter_args(%{{.*}} = %[[ARG0]]) -> (tensor<?x?xf32>) {
+//       CHECK:   scf.for %[[J:.*]]
+//       CHECK:     %[[ST:.*]] = subtensor %[[ARG0]]
+//       CHECK:     %[[ST_FILL:.*]] = linalg.fill(%[[ST]], %[[C0]]) {__internal_linalg_transform__ = "after_out_fusion_producer"} : tensor<?x?xf32>, f32 -> tensor<?x?xf32>
+//       CHECK:     %[[ST_MM_RES:.*]] = scf.for %[[K:.*]]{{.*}}iter_args(%[[BB:.*]] = %[[ST_FILL]]) -> (tensor<?x?xf32>) {
+//   CHECK-NOT:       fill
+//       CHECK:       %[[ST_MM:.*]] = linalg.matmul {__internal_linalg_transform__ = "after_out_fusion"} ins(%{{.*}}, %{{.*}} : tensor<?x?xf32>, tensor<?x?xf32>) outs(%[[BB]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+//       CHECK:       scf.yield %[[ST_MM]] : tensor<?x?xf32>
+//       CHECK:     %[[MM:.*]] = subtensor_insert %[[ST_MM_RES]] into {{.*}}
+//       CHECK:     scf.yield %[[MM]] : tensor<?x?xf32>

mlir/test/lib/Transforms/TestLinalgFusionTransforms.cpp

patterns.add<LinalgTileAndFusePattern<MatmulOp>>(
           Identifier::get("after_lhs_fusion_producer", context)),
       LinalgTransformationFilter(
           ArrayRef<Identifier>(),
           Identifier::get("after_lhs_fusion_original", context)));
 
   patterns.add<LinalgTileAndFusePattern<MatmulOp>>(
       context, dependenceGraph,
       LinalgTilingOptions().setTileSizes({32, 64, 16}).setLoopType(LoopType),
+      LinalgFusionOptions().setIndicesToFuse({2}),
+      LinalgTransformationFilter(Identifier::get("out_fusion", context),
+                                 Identifier::get("after_out_fusion", context)),
+      LinalgTransformationFilter(
+          ArrayRef<Identifier>(),
+          Identifier::get("after_out_fusion_producer", context)),
+      LinalgTransformationFilter(
+          ArrayRef<Identifier>(),
+          Identifier::get("after_out_fusion_original", context)));
+
+  patterns.add<LinalgTileAndFusePattern<MatmulOp>>(
+      context, dependenceGraph,
+      LinalgTilingOptions().setTileSizes({32, 64, 16}).setLoopType(LoopType),
       LinalgFusionOptions().setIndicesToFuse({1}),
       LinalgTransformationFilter(Identifier::get("rhs_fusion", context),
                                  Identifier::get("after_rhs_fusion", context)),
       LinalgTransformationFilter(
           ArrayRef<Identifier>(),
           Identifier::get("after_rhs_fusion_producer", context)),
       LinalgTransformationFilter(
           ArrayRef<Identifier>(),