[mlir][Linalg] Extend fusion to support WAW atm on buffers.

[hanchung]

Summary

The RAW fusion happens only if the produecer block dominates the consumer block.
The WAW pattern also works with the precondition. I.e., if a producer can
dominate the consumer, they can fairly fuse together.

Since they are all tilable, we can think the pattern like this way:

Input:

linalg_op1 view

tile_loop
  subview_2
  linalg_op2 subview_2

Tile the first Linalg op as same as the second Linalg.

tile_loop
  subview_1
  linalg_op1 subview_1

tile_loop
  subview_2
  liangl_op2 subview_2

Since the first Linalg op is tilable in the same way and the computation are
independently, it's fair to fuse it with the second Linalg op.

tile_loop
  subview_1
  linalg_op1 subview_1
  linalg_op2 subview_2

In short, this patch includes:

• Handling both RAW and WAW pattern.
• Adding a interface method to get input and output buffers.
• Exposing a method to get a StringRef of a dependency type.
• Fixing existing WAW tests and add one more use case: initialize the buffer before conv op.

[hanchung]

Revert accidental change.

[hanchung]

Fix formatting.

[mravishankar]

LGTM. Please wait for @nicolasvasilache to accept.

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td

105

Nit : rename to getInputAndOutputBuffers

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

335

Is WRW a typo?

338

You could consider changing fuseProducerOp to take an extra argument ArrayRef<LinalgDependenceGraph::DependenceType> deps to allow clients to control which deps to fuse. You could have a default of RAW and WAW dependences.

[nicolasvasilache]

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td

105

This should be either getInputOrOutputBuffer or just getBuffer.
It returns one single buffer at position $i.

mlir/test/Dialect/Linalg/fusion.mlir

336

This comment does not make sense to me.
The original comment used C and E as a shortcut for the producer of C and the producer of E.

How about:
fuse the producer of E (WAW) then the producer of C (WAR) ?

Also the resulting fusion is most likely a bug.
It's not your fault and independent from this CL.
I'll fix in a followup.

[hanchung]

Address comments.

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td

105

Yes, let's use getBuffer. It's simpler.

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

338

It needs more change to expose LinalgDependenceGraph::DependenceType because it is a enum inside LinalgDependenceGraph.

In-class enum forward declaration is not possible, according to decl.enum.

btw, I don't think a case that user would like to control the deps to fuse? Let's make the change when there is a need.

mlir/test/Dialect/Linalg/fusion.mlir

336

Oh, I misread the result. Thank you to point it out. :)

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

struct LinalgDependenceGraphElem {
     Value indexingView;
   };
   using LinalgDependences = SmallVector<LinalgDependenceGraphElem, 8>;
   using DependenceGraph = DenseMap<Operation *, LinalgDependences>;
   using dependence_iterator = LinalgDependences::const_iterator;
   using dependence_range = iterator_range<dependence_iterator>;
 
   enum DependenceType { RAR = 0, RAW, WAR, WAW, NumTypes };
+  static StringRef getDependenceTypeStr(DependenceType depType);
 
   // Builds a linalg dependence graph for the ops of type LinalgOp under `f`.
   static LinalgDependenceGraph buildDependenceGraph(Aliases &aliases, FuncOp f);
   LinalgDependenceGraph(Aliases &aliases, ArrayRef<Operation *> ops);
 
   /// Returns the X such that op -> X is a dependence of type dt.
   dependence_range getDependencesFrom(Operation *src, DependenceType dt) const;
   dependence_range getDependencesFrom(LinalgOp src, DependenceType dt) const;

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td

InterfaceMethod<
       "Return the output buffers (operands) from the current operation.",
       "Operation::operand_range", "getOutputBuffers"
     >,
 
     //===------------------------------------------------------------------===//
     // Input and Output arguments handling.
     //===------------------------------------------------------------------===//
     InterfaceMethod<
+      "Return the input or output buffer at the given index.",
+      "Value", "getInputsAndOutputsBuffer", (ins "unsigned":$i)
+    >,
+    InterfaceMethod<
       "Return the number of inputs and outputs, irrespective of their buffer "
       "or tensor type.",
       "unsigned", "getNumInputsAndOutputs"
     >,
     InterfaceMethod<
       "Return the number of inputs, irrespective of their buffer or tensor "
       "type, and output buffers",
       "unsigned", "getNumInputsAndOutputBuffers"

mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h

Operation::operand_range getOutputBuffers() {
     auto range = this->getOperation()->getOperands();
     return {range.begin() + nInputs(),
             range.begin() + getNumInputsAndOutputBuffers()};
   }
 
   //==========================================================================//
   // Input and Output arguments handling.
   //==========================================================================//
+  Value getInputsAndOutputsBuffer(unsigned i) {
+    assert(i < getNumInputsAndOutputBuffers() &&
+           "overflowing inputs and outputs buffer index");
+    return this->getOperation()->getOperand(i);
+  }
   /// Return the number of inputs and outputs, irrespective of their buffer or
   /// tensor type.
   unsigned getNumInputsAndOutputs() { return nInputs() + nOutputs(); }
   /// Return the number of inputs, irrespective of their buffer or tensor type,
   /// and output buffers.
   unsigned getNumInputsAndOutputBuffers() {
     assert(this->getOperation()->getNumResults() <= nOutputs());
     return nInputs() + nOutputs() - this->getOperation()->getNumResults();

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

 
 #define DEBUG_TYPE "linalg-dependence-analysis"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
 using llvm::dbgs;
 
-#ifndef NDEBUG
-static StringRef toStringRef(LinalgDependenceGraph::DependenceType dt) {
-  switch (dt) {
-  case LinalgDependenceGraph::DependenceType::RAW:
-    return "RAW";
-  case LinalgDependenceGraph::DependenceType::RAR:
-    return "RAR";
-  case LinalgDependenceGraph::DependenceType::WAR:
-    return "WAR";
-  case LinalgDependenceGraph::DependenceType::WAW:
-    return "WAW";
-  default:
-    break;
-  }
-  llvm_unreachable("Unexpected DependenceType");
-}
-#endif
-
 Value Aliases::find(Value v) {
   if (v.isa<BlockArgument>())
     return v;
 
   auto it = aliases.find(v);
   if (it != aliases.end()) {
     assert(it->getSecond().getType().isa<MemRefType>() && "Memref expected");
     return it->getSecond();
if (auto view = dyn_cast_or_null<SubViewOp>(v.getDefiningOp())) {
       v = view.source();
       continue;
     }
     llvm::errs() << "View alias analysis reduces to: " << v << "\n";
     llvm_unreachable("unsupported view alias case");
   }
 }
 
+StringRef LinalgDependenceGraph::getDependenceTypeStr(DependenceType depType) {
+  switch (depType) {
+  case LinalgDependenceGraph::DependenceType::RAW:
+    return "RAW";
+  case LinalgDependenceGraph::DependenceType::RAR:
+    return "RAR";
+  case LinalgDependenceGraph::DependenceType::WAR:
+    return "WAR";
+  case LinalgDependenceGraph::DependenceType::WAW:
+    return "WAW";
+  default:
+    break;
+  }
+  llvm_unreachable("Unexpected DependenceType");
+}
+
 LinalgDependenceGraph
 LinalgDependenceGraph::buildDependenceGraph(Aliases &aliases, FuncOp f) {
   SmallVector<Operation *, 8> linalgOps;
   f.walk([&](LinalgOp op) { linalgOps.push_back(op); });
   return LinalgDependenceGraph(aliases, linalgOps);
 }
 
 LinalgDependenceGraph::LinalgDependenceGraph(Aliases &aliases,
                                              ArrayRef<Operation *> ops)
     : aliases(aliases), linalgOps(ops.begin(), ops.end()) {
   for (auto en : llvm::enumerate(linalgOps)) {
     assert(isa<LinalgOp>(en.value()) && "Expected value for LinalgOp");
     linalgOpPositions.insert(std::make_pair(en.value(), en.index()));
   }
   for (unsigned i = 0, e = ops.size(); i < e; ++i) {
     for (unsigned j = i + 1; j < e; ++j) {
       addDependencesBetween(cast<LinalgOp>(ops[i]), cast<LinalgOp>(ops[j]));
     }
   }
 }
 
 void LinalgDependenceGraph::addDependenceElem(DependenceType dt,
                                               LinalgOpView indexingOpView,
                                               LinalgOpView dependentOpView) {
-  LLVM_DEBUG(dbgs() << "\nAdd dep type " << toStringRef(dt) << ":\t"
+  LLVM_DEBUG(dbgs() << "\nAdd dep type " << getDependenceTypeStr(dt) << ":\t"
                     << *indexingOpView.op << " -> " << *dependentOpView.op);
   dependencesFromGraphs[dt][indexingOpView.op].push_back(
       LinalgDependenceGraphElem{dependentOpView, indexingOpView.view});
   dependencesIntoGraphs[dt][dependentOpView.op].push_back(
       LinalgDependenceGraphElem{indexingOpView, dependentOpView.view});
 }
 
 LinalgDependenceGraph::dependence_range
for (auto dependence : getDependencesFrom(src, dt)) {
       auto interimPos = linalgOpPositions.lookup(dependence.dependentOpView.op);
       // Skip if not interleaved.
       if (interimPos >= dstPos || interimPos <= srcPos)
         continue;
       if (view && !aliases.alias(view, dependence.indexingView))
         continue;
       auto *op = dependence.dependentOpView.op;
       LLVM_DEBUG(dbgs() << "\n***Found covering dependence of type "
-                        << toStringRef(dt) << ": " << *src << " -> " << *op
-                        << " on " << dependence.indexingView);
+                        << getDependenceTypeStr(dt) << ": " << *src << " -> "
+                        << *op << " on " << dependence.indexingView);
       res.push_back(op);
     }
   }
   return res;
 }

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

static LinalgOp fuse(Value producedView, LinalgOp producer, LinalgOp consumer,
   }
   if (auto convOp = dyn_cast<linalg::ConvOp>(consumer.getOperation())) {
     // TODO(ntv): add a level of indirection to linalg.generic.
     if (convOp.padding())
       llvm_unreachable("Unexpected conv with padding");
   }
 
   auto subView = dyn_cast_or_null<SubViewOp>(
-      consumer.getInput(consumerIdx).getDefiningOp());
-  auto slice =
-      dyn_cast_or_null<SliceOp>(consumer.getInput(consumerIdx).getDefiningOp());
+      consumer.getInputsAndOutputsBuffer(consumerIdx).getDefiningOp());
+  auto slice = dyn_cast_or_null<SliceOp>(
+      consumer.getInputsAndOutputsBuffer(consumerIdx).getDefiningOp());
   assert(subView || slice);
   (void)subView;
   (void)slice;
 
   // loopToOperandRangesMaps are permutations-only by construction:
   //   we can always identify a data dimension with a (at least one) loop
   //   dimension.
   AffineMap producerMap =
bool mlir::linalg::isFusableInto(const LinalgDependenceGraph &graph,
   if (!graph.findCoveringDependences(producer, consumer).empty()) {
     LLVM_DEBUG(dbgs() << "\n***Not fusable due to an interleaved dependence:\t"
                       << *producer.getOperation());
     return false;
   }
   return true;
 }
 
-// Only consider RAW atm.
-Optional<FusionInfo> mlir::linalg::fuseProducerOf(
-    OpBuilder &b, LinalgOp consumer, unsigned consumerIdx,
-    const LinalgDependenceGraph &graph, OperationFolder *folder) {
+static Optional<FusionInfo>
+fuseProducerOfDep(OpBuilder &b, LinalgOp consumer, unsigned consumerIdx,
+                  const LinalgDependenceGraph &graph, OperationFolder *folder,
+                  LinalgDependenceGraph::DependenceType depType) {
   assert(consumer.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
   LLVM_DEBUG(dbgs() << "\nStart examining consumer: "
                     << *consumer.getOperation());
-  for (auto dependence : graph.getDependencesInto(
-           consumer, LinalgDependenceGraph::DependenceType::RAW)) {
+  for (auto dependence : graph.getDependencesInto(consumer, depType)) {
     LLVM_DEBUG(dbgs() << "\n***Consider producer:\t"
                       << *dependence.dependentOpView.op << "\n");
     auto producer = cast<LinalgOp>(dependence.dependentOpView.op);
     if (isa<linalg::IndexedGenericOp>(dependence.dependentOpView.op)) {
       LLVM_DEBUG(dbgs() << "Not fusing indexed_generic producer");
       continue;
     }
 
     // Check that the dependence is indeed on the input `consumerIdx` view.
     auto consumedView = dependence.indexingView;
-    if (consumer.getInput(consumerIdx) != consumedView)
+    if (consumer.getInputsAndOutputsBuffer(consumerIdx) != consumedView)
       continue;
 
     // Consumer consumes this view, `isStructurallyFusableProducer` also checks
     // whether it is a strict subview of the producer view.
     auto producedView = dependence.dependentOpView.view;
     auto producerIdx = producer.getIndexOfOutputBuffer(producedView).getValue();
     // `consumerIdx` and `producerIdx` exist by construction.
-    LLVM_DEBUG(dbgs() << "\nRAW producer: " << *producer.getOperation()
-                      << " view: " << producedView
-                      << " output index: " << producerIdx);
+    LLVM_DEBUG(dbgs() << "\n"
+                      << LinalgDependenceGraph::getDependenceTypeStr(depType)
+                      << "producer: " << *producer.getOperation() << " view: "
+                      << producedView << " output index: " << producerIdx);
 
     // Must be a subview or a slice to guarantee there are loops we can fuse
     // into.
     auto subView = dyn_cast_or_null<SubViewOp>(consumedView.getDefiningOp());
     auto slice = dyn_cast_or_null<SliceOp>(consumedView.getDefiningOp());
     if (!subView && !slice) {
       LLVM_DEBUG(dbgs() << "\nNot fusable (not a subview or slice)");
       continue;
for (auto dependence : graph.getDependencesInto(consumer, depType)) {
     auto fusedProducer = fuse(producedView, producer, consumer, consumerIdx,
                               producerIdx, folder);
 
     return FusionInfo{producer, fusedProducer};
   }
   return llvm::None;
 }
 
+// Only consider RAW and WRW atm.
+Optional<FusionInfo> mlir::linalg::fuseProducerOf(
+    OpBuilder &b, LinalgOp consumer, unsigned consumerIdx,
+    const LinalgDependenceGraph &graph, OperationFolder *folder) {
+  SmallVector<LinalgDependenceGraph::DependenceType, 4> deps = {
+      LinalgDependenceGraph::DependenceType::RAW,
+      LinalgDependenceGraph::DependenceType::WAW,
+  };
+  for (auto dep : deps) {
+    if (auto res =
+            fuseProducerOfDep(b, consumer, consumerIdx, graph, folder, dep))
+      return res;
+  }
+  return llvm::None;
+}
+
 /// Checks if two Generic ops are fusible, when one is a producer and another is
 /// a consumer (with the result of the producer being the `consumerIdx` operand
 /// of the consumer).
 static bool areTensorOpsFusible(LinalgOp producer, LinalgOp consumer,
                                 unsigned consumerIdx) {
   // Verify that the producer and consumer are ops on tensors.
   if (!producer.hasTensorSemantics() || !consumer.hasTensorSemantics())
     return false;
f.walk([&](LinalgOp op) {
     if (op.hasBufferSemantics())
       linalgOps.push_back(op);
   });
 
   // TODO(pifon, ntv): LinalgDependenceGraph should be able to update itself.
   // The current naive and expensive reconstruction of the graph should be
   // removed.
   for (auto *op : llvm::reverse(linalgOps)) {
-    for (unsigned id = 0, e = LinalgOp(op).getNumInputs(); id < e; ++id) {
+    for (unsigned id = 0, e = LinalgOp(op).getNumInputsAndOutputBuffers();
+         id < e; ++id) {
       linalg::Aliases aliases;
       linalg::LinalgDependenceGraph graph(aliases, linalgOps);
       if (auto info = fuseProducerOf(b, op, id, graph, &folder)) {
         auto *originalOp = info->originalProducer.getOperation();
         eraseSet.insert(originalOp);
         auto *originalOpInLinalgOpsVector =
             std::find(linalgOps.begin(), linalgOps.end(), originalOp);
         *originalOpInLinalgOpsVector = info->fusedProducer.getOperation();

mlir/test/Dialect/Linalg/fusion.mlir

loop.for %arg6 = %c0 to %2 step %c3 {
           memref<?x?xf32, offset: ?, strides: [?, ?]>
       }
     }
   }
   return %E : memref<?x?xf32, offset: 0, strides: [?, 1]>
 }
 // CHECK-LABEL: func @f1
 // CHECK:   (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}})
-// No RAW dependences, the pass does not fuse RAR atm.
-// CHECK: linalg.matmul
 // CHECK: loop.for
 // CHECK:   loop.for
 // CHECK:     loop.for
 // CHECK:       linalg.matmul
+// CHECK:       linalg.matmul
 
 // -----
 
 // CHECK-DAG: #[[strided2D:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s0 + d1 * s1)>
 func @f2(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>,
          %B: memref<?x?xf32, offset: 0, strides: [?, ?]>,
          %C: memref<?x?xf32, offset: 0, strides: [?, ?]>,
          %D: memref<?x?xf32, offset: 0, strides: [?, ?]>,
loop.for %arg6 = %c0 to %2 step %c3 {
           memref<?x?xf32, offset: ?, strides: [?, ?]>
       }
     }
   }
   return %E : memref<?x?xf32, offset: 0, strides: [?, ?]>
 }
 // CHECK-LABEL: func @f6
 // CHECK:  (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}})
-// Cannot fuse C due to interleaved read of C that would be bypassed.
-// Cannot fuse E (WAW).
-// CHECK:  linalg.matmul
-// CHECK:  linalg.matmul
+// Fuse B and C, then fuse with A.
 // CHECK:  loop.for
 // CHECK:    loop.for
 // CHECK:      loop.for
 // CHECK:        linalg.matmul
-// CHECK-NOT:      linalg.matmul
+// CHECK:        linalg.matmul
+// CHECK:        linalg.matmul
 
 // -----
 
 func @f7(%A: memref<?x?xf32, offset: 0, strides: [?, ?]>,
          %B: memref<?x?xf32, offset: 0, strides: [?, ?]>,
          %C: memref<?x?xf32, offset: 0, strides: [?, ?]>,
          %D: memref<?x?xf32, offset: 0, strides: [?, ?]>,
          %E: memref<?x?xf32, offset: 0, strides: [?, ?]>

 // CHECK:       linalg.generic
 // CHECK:         linalg.yield
 // CHECK:       linalg.generic
 // CHECK:         subf
 // CHECK:         linalg.yield
 // CHECK:       linalg.generic
 // CHECK:         exp
 // CHECK:         linalg.yield
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2) -> (d0, d1 - d2)>
+#map1 = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3, s4] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3 * s4)>
+#map2 = affine_map<()[s0] -> (s0 + 3)>
+
+func @fill_and_conv(%arg0: memref<1x4x5x1xf32>, %arg1: memref<2x3x1x1xf32>, %arg2: memref<1x4x5x1xf32>) {
+  %cst = constant 0.000000e+00 : f32
+  linalg.fill(%arg2, %cst) : memref<1x4x5x1xf32>, f32
+
+  %c4 = constant 4 : index
+  %c1 = constant 1 : index
+  %c0 = constant 0 : index
+  %c2 = constant 2 : index
+  %c3 = constant 3 : index
+  %4 = dim %arg1, 0 : memref<2x3x1x1xf32>
+  %5 = dim %arg1, 1 : memref<2x3x1x1xf32>
+  %6 = dim %arg0, 0 : memref<1x4x5x1xf32>
+  %7 = dim %arg0, 1 : memref<1x4x5x1xf32>
+  %8 = dim %arg0, 3 : memref<1x4x5x1xf32>
+  %9 = dim %arg2, 0 : memref<1x4x5x1xf32>
+  %10 = dim %arg2, 1 : memref<1x4x5x1xf32>
+  %11 = dim %arg2, 2 : memref<1x4x5x1xf32>
+  %12 = dim %arg2, 3 : memref<1x4x5x1xf32>
+  %13 = linalg.range %c0 : %6 : %c2 : !linalg.range
+  %14 = linalg.range %c0 : %10 : %c3 : !linalg.range
+  loop.for %arg3 = %c0 to %6 step %c2 {
+    loop.for %arg4 = %c0 to %10 step %c3 {
+      %15 = affine.min #map0(%c2, %c1, %arg3)
+      %16 = affine.apply #map2()[%7]
+      %17 = affine.min #map0(%16, %c4, %arg4)
+      %18 = dim %arg0, 2 : memref<1x4x5x1xf32>
+      %19 = dim %arg0, 3 : memref<1x4x5x1xf32>
+      %20 = subview %arg0[%arg3, %arg4, %c0, %c0] [%15, %17, %18, %19] [%c1, %c1, %c1, %c1] : memref<1x4x5x1xf32> to memref<?x?x?x?xf32, #map1>
+      %21 = affine.min #map0(%c2, %c1, %arg3)
+      %22 = affine.min #map0(%c3, %c4, %arg4)
+      %23 = dim %arg2, 2 : memref<1x4x5x1xf32>
+      %24 = dim %arg2, 3 : memref<1x4x5x1xf32>
+      %25 = subview %arg2[%arg3, %arg4, %c0, %c0] [%21, %22, %23, %24] [%c1, %c1, %c1, %c1] : memref<1x4x5x1xf32> to memref<?x?x?x?xf32, #map1>
+      linalg.conv(%arg1, %20, %25) {dilations = [1, 1], strides = [1, 1]} : memref<2x3x1x1xf32>, memref<?x?x?x?xf32, #map1>, memref<?x?x?x?xf32, #map1>
+    }
+  }
+  return
+}
+// CHECK-LABEL: func @fill_and_conv
+// CHECK: loop.for
+// CHECK:   loop.for
+// CHECK:     linalg.fill
+// CHECK:     linalg.conv