[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

if (isa<GenericOp>(producer) || isa<IndexedGenericOp>(producer)) {
       if (linalgOpProducer.hasTensorSemantics())
         return FuseGenericOpsOnTensors::fuse(linalgOpProducer, linalgOpConsumer,
                                              consumerIdx, rewriter, folder);
     } else if (auto reshapeOpProducer = dyn_cast<TensorReshapeOp>(producer)) {
       if (auto genericOpConsumer = dyn_cast<GenericOp>(consumer)) {
         return FuseTensorReshapeOpAsProducer<GenericOp>::fuse(
             reshapeOpProducer, genericOpConsumer, consumerIdx, rewriter,
             folder);
+      } else if (auto indexedGenericOpConsumer =
+                     dyn_cast<IndexedGenericOp>(consumer)) {
+        return FuseTensorReshapeOpAsProducer<IndexedGenericOp>::fuse(
+            reshapeOpProducer, indexedGenericOpConsumer, consumerIdx, rewriter,
+            folder);
       }
     } else if (auto constantOpProducer = dyn_cast<ConstantOp>(producer)) {
       if (auto genericOpConsumer = dyn_cast<GenericOp>(consumer)) {
         return FuseConstantOpAsProducer<GenericOp>::fuse(
             constantOpProducer, genericOpConsumer, consumerIdx, rewriter,
             folder);
       }
     }
     return nullptr;
   }
 
   // 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;
   }
 
   return nullptr;
 }
 
 namespace {

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