[mlir][Linalg] Make sure the value used to obtain tile shape is a SubViewOp/SubTensorOp.

[mravishankar]

Summary

Current logic used to get the bounds of loop depends on the use of
getOrCreateRange method on SubViewOp and SubTensorOp. Make sure
that the value/dim used to compute the range is from such ops.
This fix is a reasonable WAR, but a btter fix would be to make
getOrCreateRange method be a method of ViewInterface.

[mravishankar]

Landed as part of https://reviews.llvm.org/rGe65a5e5b00a37700a79e0a9f2fb1c1e60a2584bf . Closing.

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

struct ShapeDimension {
   unsigned dimension;
 };
 
 // Given an `op`, returns the first (`shape`, `dimension`) pair that identifies
 // the loop range at `loopDepth`. The semantics of the loopToOperandRangesMaps
 // guarantees at least one such dimension is found. If multiple candidates exist
 // they must agree by construction (i.e. have the same size) and we just return
 // the first one.
-static ShapeDimension getShapeDefiningLoopRange(LinalgOp op,
-                                                unsigned loopDepth) {
+static ShapeDimension
+getShapeDefiningLoopRange(LinalgOp op, unsigned loopDepth,
+                          bool fromSubViewOpOnly = false) {
   auto maps = op.indexing_maps();
   // Iterate over the inputs and outputs in order.
   // Extract the subranges from the linearized ranges.
   SmallVector<Value, 8> ios(op.getInputsAndOutputBuffers());
   for (auto en : llvm::enumerate(ios)) {
+    // The method `getRangeFromOperandShape` requires using SubViewOp or
+    // SubTensorOps. If the value isnt defined from there continue.
+    // todo: The method should be adapted to get the values from
+    // `ViewInterface`. The interface needs a `getOrCreateRanges` method which
+    // currently returns a `linalg.range`. The fix here is to move this op to
+    // `std` dialect and add the method to `ViewInterface`.
+    if (fromSubViewOpOnly &&
+        !isa_and_nonnull<SubViewOp, SubTensorOp>(en.value().getDefiningOp()))
+      continue;
+
     unsigned idx = en.index();
     auto map = maps[idx].cast<AffineMapAttr>().getValue();
     LLVM_DEBUG(llvm::dbgs()
                << "getShapeDefiningLoopRange I/O idx: " << idx << "\n");
     LLVM_DEBUG(llvm::dbgs()
                << "getShapeDefiningLoopRange map: " << map << "\n");
     Value shape = en.value();
     SmallVector<Value, 8> shapeRanges(map.getNumResults(), nullptr);

 fuseOperations(OpBuilder &builder, LinalgOp tiledOp,
                ArrayRef<LinalgOp> fusionCandidates,
                const FusableOpDependencesTy &fusableDependences,
                const std::set<unsigned> &fusedLoops) {
   OpBuilder::InsertionGuard guard(builder);
   builder.setInsertionPoint(tiledOp);
   DenseMap<unsigned, Range> fusedLoopsAndRanges;
   for (unsigned loop : fusedLoops) {
-    ShapeDimension shapeDim = getShapeDefiningLoopRange(tiledOp, loop);
+    ShapeDimension shapeDim = getShapeDefiningLoopRange(tiledOp, loop, true);
     fusedLoopsAndRanges[loop] = getRangeFromOperandShape(
         builder, tiledOp.getLoc(), shapeDim.shape, shapeDim.dimension);
   }
   SmallVector<LinalgOp, 1> fusedOps(fusionCandidates.size());
   for (auto candidate : enumerate(llvm::reverse(fusionCandidates))) {
     LinalgOp fusedOp = fuse(builder, candidate.value(), fusedLoopsAndRanges);
     fusedOps[fusionCandidates.size() - candidate.index() - 1] = fusedOp;
     builder.setInsertionPoint(fusedOp);

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

// expected-remark @+1 {{unhandled fusion of ops in different basic blocks}}
           ins(%6, %7 : memref<?x?xf32, #map3>, memref<?x?xf32, #map3>)
           outs(%8 : memref<?x?xf32, #map3>)
       }
       scf.yield
     }
     return
   }
 }
+
+// -----
+
+module {
+  func @basic_conv_fusion(%arg0: memref<?x?x?x?xf32>, %arg1: memref<?x?x?x?xf32>,
+                          %arg2: memref<?x?x?x?xf32>) {
+    %cst = constant 0.000000e+00 : f32
+    linalg.fill(%arg2, %cst) : memref<?x?x?x?xf32>, f32
+    linalg.conv(%arg0, %arg1, %arg2) {
+      dilations = [1, 1], strides = [1, 1],
+      __internal_linalg_transform__ = "basic_fusion"} :
+      memref<?x?x?x?xf32>, memref<?x?x?x?xf32>, memref<?x?x?x?xf32>
+    return
+  }
+}
+//      CHECK: func @basic_conv_fusion
+//      CHECK:   linalg.fill
+// CHECK-SAME:     __internal_linalg_transform__ = "after_basic_fusion_original"
+//      CHECK:  scf.parallel (%{{.+}}, %{{.+}}, %{{.+}})
+// CHECK-SAME:  {
+//      CHECK:    linalg.fill
+// CHECK-SAME:      __internal_linalg_transform__ = "after_basic_fusion_producer"
+//      CHECK:    linalg.conv
+// CHECK-SAME:      __internal_linalg_transform__ = "after_basic_fusion"
+//      CHECK:  }
+//      CHECK:  linalg.conv
+// CHECK-SAME:    __internal_linalg_transform__ = "after_basic_fusion_original"

mlir/test/lib/Transforms/TestLinalgFusionTransforms.cpp

struct TestLinalgFusionTransforms
 
   void runOnFunction() override;
 };
 } // namespace
 
 static void fillFusionPatterns(MLIRContext *context,
                                const LinalgDependenceGraph &dependenceGraph,
                                OwningRewritePatternList &patterns) {
-  patterns.insert<LinalgTileAndFusePattern<MatmulOp>>(
+  patterns.insert<LinalgTileAndFusePattern<MatmulOp>,
+                  LinalgTileAndFusePattern<ConvOp>>(
       context, dependenceGraph,
       LinalgTilingOptions()
           .setTileSizes({32, 64, 16})
           .setLoopType(LinalgTilingLoopType::ParallelLoops),
       LinalgFusionOptions().setIndicesToFuse({2}),
       LinalgMarker(Identifier::get("basic_fusion", context),
                    Identifier::get("after_basic_fusion", context)),
       LinalgMarker(ArrayRef<Identifier>(),