[mlir][Linalg] Add support for fusion between indexed_generic ops and tensor_reshape ops

[hanchung]

Summary

The fusion for tensor_reshape is embedding the information to indexing maps,
thus the exising pattenr also works for indexed_generic ops.

Depends On D80347

[mravishankar]

Good to know that this minimal change works as expected (I hoped it would). Thanks for adding tests.
Based on the previous patch in this stack, if we can combine the methods for fusion of generic/indexed-generic producer/consumer then I think this change will not be needed. The tests are still useful though. Will wait to review this again after that.

[hanchung]

Update parents

[hanchung]

In D80348#2052290, @mravishankar wrote:

Good to know that this minimal change works as expected (I hoped it would). Thanks for adding tests.
Based on the previous patch in this stack, if we can combine the methods for fusion of generic/indexed-generic producer/consumer then I think this change will not be needed. The tests are still useful though. Will wait to review this again after that.

This is a bit tricky, we still need the change to deduce the type of the op to create (whether this is GenericOp or IndexedGenericOp).

[mravishankar]

This is a bit tricky, we still need the change to deduce the type of the op to create (whether this is GenericOp or IndexedGenericOp).

I agree. This is more evidence for me that the indexed-generic and generic should really become op-interfaces after the named ops are replaced with generated ops in Linalg. @nicolasvasilache what do you think. I can start a discussion on discourse if thats a better forum.

[hanchung]

rebase

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

Operation *mlir::linalg::fuseTensorOps(PatternRewriter &rewriter,
   }
 
   // Fuse when consumer is a TensorReshapeOp.
   if (TensorReshapeOp reshapeOp = dyn_cast<TensorReshapeOp>(consumer)) {
     if (auto genericOpProducer = dyn_cast<GenericOp>(producer)) {
       if (genericOpProducer.hasTensorSemantics())
         return FuseTensorReshapeOpAsConsumer<GenericOp>::fuse(
             genericOpProducer, reshapeOp, consumerIdx, rewriter, folder);
+    } else if (auto indexedGenericOpProducer =
+                   dyn_cast<IndexedGenericOp>(producer)) {
+      if (indexedGenericOpProducer.hasTensorSemantics())
+        return FuseTensorReshapeOpAsConsumer<IndexedGenericOp>::fuse(
+            indexedGenericOpProducer, reshapeOp, consumerIdx, rewriter, folder);
     }
     return nullptr;
   }
 
   // Fuse when consumer is IndexedGenericOp.
   if (IndexedGenericOp indexedGenericOp =
           dyn_cast<IndexedGenericOp>(consumer)) {
     if (!indexedGenericOp.hasTensorSemantics())
       return nullptr;
     if (auto genericOpProducer = dyn_cast<GenericOp>(producer)) {
       if (genericOpProducer.hasTensorSemantics())
         return FuseIndexedGenericOpsOnTensors::fuse(
             genericOpProducer, indexedGenericOp, consumerIdx, rewriter, folder);
     }
     if (auto indexedGenericOpProducer = dyn_cast<IndexedGenericOp>(producer)) {
-      if (indexedGenericOpProducer.hasTensorSemantics())
+      if (indexedGenericOpProducer.hasTensorSemantics()) {
         return FuseIndexedGenericOpsOnTensors::fuse(
             indexedGenericOpProducer, indexedGenericOp, consumerIdx, rewriter,
             folder);
       }
+    } else if (auto reshapeOpProducer = dyn_cast<TensorReshapeOp>(producer)) {
+      return FuseTensorReshapeOpAsProducer<IndexedGenericOp>::fuse(
+          reshapeOpProducer, indexedGenericOp, consumerIdx, rewriter, folder);
+    }
     return nullptr;
   }
-
   return nullptr;
 }
 
 namespace {
 /// Patterns to fuse a generic op, with the producer of its operands.
 template <typename LinalgOpTy>
 struct FuseTensorOps : public OpRewritePattern<LinalgOpTy> {
   using OpRewritePattern<LinalgOpTy>::OpRewritePattern;

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

 //      CHECK:   %[[VAL1:.+]] = addi %[[ARG2]], %[[ADD_OPERAND1]] : i32
 //      CHECK:   %[[VAL2:.+]] = subi %[[VAL1]], %[[SUB_OPERAND1]] : i32
 //      CHECK:   %[[ADD_OPERAND2:.+]] = index_cast %[[ARG0]] : index to i32
 //      CHECK:   %[[SUB_OPERAND2:.+]] = index_cast %[[ARG1]] : index to i32
 //      CHECK:   %[[VAL3:.+]] = addi %[[VAL2]], %[[ADD_OPERAND2]] : i32
 //      CHECK:   %[[VAL4:.+]] = subi %[[VAL3]], %[[SUB_OPERAND2]] : i32
 //      CHECK:   linalg.yield %[[VAL4]] : i32
 //   CHECK-NOT: linalg.indexed_generic
+
+// -----
+
+// CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 * 4 + d2, d3)>
+// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+func @indexed_generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x?xi32>)
+  -> tensor<?x?x4x?xi32> {
+  %0 = linalg.tensor_reshape %arg0 [affine_map<(i, j, k, l) -> (i)>,
+                                    affine_map<(i, j, k, l) -> (j, k)>,
+                                    affine_map<(i, j, k, l) -> (l)>] :
+    tensor<?x?x?xi32> into tensor<?x?x4x?xi32>
+  %1 = linalg.indexed_generic {
+    args_in = 1 : i64,
+    args_out = 1 : i64,
+    indexing_maps = [#map0, #map0],
+    iterator_types = ["parallel", "parallel", "parallel", "parallel"] } %0 {
+  ^bb0(%arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: i32):       // no predecessors
+    %2 = index_cast %arg2 : index to i32
+    %3 = addi %arg6, %2 : i32
+    linalg.yield %3 : i32
+  }: tensor<?x?x4x?xi32> -> tensor<?x?x4x?xi32>
+  return %1 : tensor<?x?x4x?xi32>
+}
+
+// CHECK-LABEL: func @indexed_generic_op_reshape_producer_fusion
+//   CHECK-NOT: linalg.tensor_reshape
+//       CHECK: linalg.indexed_generic
+//  CHECK-SAME:   args_in = 1
+//  CHECK-SAME:   args_out = 1
+//  CHECK-SAME:   indexing_maps = [#[[MAP0]], #[[MAP1]]]
+//   CHECK-NOT: linalg.tensor_reshape
+
+// -----
+
+// CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 * 20 + d2 * 5 + d3)>
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+func @indexed_generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?x4x5xi32>)
+  -> tensor<?x?xi32> {
+  %0 = linalg.indexed_generic {
+    args_in = 1 : i64,
+    args_out = 1 : i64,
+    indexing_maps = [#map0, #map0],
+    iterator_types = ["parallel", "parallel", "parallel", "parallel"] } %arg0 {
+  ^bb0(%arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: i32):       // no predecessors
+    %2 = index_cast %arg2 : index to i32
+    %3 = addi %arg6, %2 : i32
+    linalg.yield %3 : i32
+  }: tensor<?x?x4x5xi32> -> tensor<?x?x4x5xi32>
+  %1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>,
+                                 affine_map<(i, j, k, l) -> (j, k, l)>] :
+    tensor<?x?x4x5xi32> into tensor<?x?xi32>
+  return %1 : tensor<?x?xi32>
+}
+
+// CHECK-LABEL: func @indexed_generic_op_reshape_consumer_fusion
+//   CHECK-NOT: linalg.tensor_reshape
+//       CHECK: linalg.indexed_generic
+//  CHECK-SAME:   args_in = 1
+//  CHECK-SAME:   args_out = 1
+//  CHECK-SAME:   indexing_maps = [#[[MAP0]], #[[MAP1]]]
+//   CHECK-NOT: linalg.tensor_reshape