Relax FuseTensorReshapeOpAsproducer identity mapping constraint

[asaadaldien]

Summary

Allow fusing linalg.tensor_reshape with its consumer when the consumer is applying a permutation map (e.g linalg.generic implementing a transpose)

[hanchung]

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

328

I think isPermutation covers isIdentity, so maybe no need to check the first case, ie, useIndexMap.isIdentity() || ?

385–387

I think we actually just want to make sure the map is invertible, so the comment is a bit not accurate. The identity map or permutation map is guaranteed to be invertible, so you don't need to check if invMap is null.

[asaadaldien]

comments...

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

328

Checked, isPermutation doesn't exclude identity.

385–387

The comment above implies consumer map fusedIndexMaps[consumerIdx] is invertible it's just phrased differently.

[mravishankar]

Thanks Ahmed, just one minor comment.

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

328

Do we need to make a distinction for when it is asProducer and not. I think you can make it isPermutation always? Then you can maybe also add a test where the reshape is the consumer?

[asaadaldien]

In D88869#2314801, @mravishankar wrote:

Thanks Ahmed, just one minor comment.

In D88869#2314801, @mravishankar wrote:

Thanks Ahmed, just one minor comment.

I need to look closely into FuseTensorReshapeOpAsConsumer to see why it isn't producing the correct indexing map for the permutation case. I can do this in a follow up diff.

NVM, that was a silly test bug :D, ignore my comment.

[asaadaldien]

enable premutations for tensor_reshape as consumer

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

static bool isTensorReshapeOpFusible(TensorReshapeOp reshapeOp,
   // Reshape is fusible with its consumer (i.e. reshape as a producer) when its
   // operand is of lesser rank than the result. Fusing when operand has higher
   // rank will require use of mods and divs in the indexing maps of the fused op
   // which would make it non-invertible. Similarly reshape is fused with its
   // producer (i.e. reshape as consumer) only if the return type has lesser
   // rank.
   if ((asProducer && returnType.getRank() < operandType.getRank()) ||
       (!asProducer && operandType.getRank() < returnType.getRank()))
     return false;
-  return useIndexMap.isIdentity();
+  return asProducer ? useIndexMap.isIdentity() || useIndexMap.isPermutation()
+                    : useIndexMap.isIdentity();
 }
 
 /// Based on the type of `op` create a linalg op of the same type, i.e. if `op`
 /// is a linalg.generic operation, the create a `linalg.generic` operation with
 /// the given `args`. Expects `op` to be `linalg.generic` or
 /// `linalg.indexed_generic`.
 template <typename... Args>
 static LinalgOp createLinalgOpOfSameType(LinalgOp op, PatternRewriter &rewriter,
static LinalgOp fuse(TensorReshapeOp producer, LinalgOp consumer,
     // Compute indexing_maps for the fused operation. The indexing_maps for the
     // operands of the consumers that arent fused are the same.
     SmallVector<AffineMap, 4> fusedIndexMaps =
         llvm::to_vector<4>(llvm::map_range(
             consumer.indexing_maps(), [](Attribute attr) -> AffineMap {
               return attr.cast<AffineMapAttr>().getValue();
             }));
 
+    // Accepted consumer maps are either identity or permutation.
+    auto invMap = inversePermutation(fusedIndexMaps[consumerIdx]);
+
     // Compute the indexing map to use for the operand of the producer.
-    AffineMap modifiedMap = linearizeCollapsedDims(
-        fusedIndexMaps[consumerIdx], producer.getResultType().getShape(),
+    AffineMap modifiedMap =
+        linearizeCollapsedDims(invMap, producer.getResultType().getShape(),
                                producer.getReassociationMaps());
     for (AffineExpr expr : modifiedMap.getResults()) {
       if (!expr.isPureAffine())
         return nullptr;
     }
     fusedIndexMaps[consumerIdx] = modifiedMap;
 
     // Further check that the resulting index maps can be fused and
     // inverted. Without this the resultant op is not legal.

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

func @indexed_generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?x4x5xi32>)
   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:   indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
 //   CHECK-NOT: linalg.tensor_reshape
+
+// -----
+
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1 + d2 * 7)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+
+#map0 = affine_map<(d0, d1, d2) -> (d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+func @generic_op_021_permultation_reshape_producer_fusion(%arg0 : tensor<3x35xf32>) -> tensor<3x7x5xf32> {
+  %0 = linalg.tensor_reshape %arg0 [#map0, #map1] : tensor<3x35xf32> into tensor<3x5x7xf32>
+  %1 = linalg.generic {indexing_maps = [#map2, #map3], iterator_types = ["parallel", "parallel", "parallel"]} ins(%0 : tensor<3x5x7xf32>) {
+    ^bb0(%arg2: f32):  // no predecessors
+      linalg.yield %arg2 : f32
+    } -> tensor<3x7x5xf32>
+    return %1 : tensor<3x7x5xf32>
+}
+
+// CHECK-LABEL: func @generic_op_021_permultation_reshape_producer_fusion
+//   CHECK-NOT: linalg.tensor_reshape
+//       CHECK: linalg.generic
+//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
+//   CHECK-NOT: linalg.tensor_reshape
+
+// -----
+
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d2, d0 * 7 + d1)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+
+#map0 = affine_map<(d0, d1, d2) -> (d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+func @generic_op_120_permultation_reshape_producer_fusion(%arg0 : tensor<3x35xf32>) -> tensor<5x7x3xf32> {
+  %0 = linalg.tensor_reshape %arg0 [#map0, #map1] : tensor<3x35xf32> into tensor<3x5x7xf32>
+  %1 = linalg.generic {indexing_maps = [#map2, #map3], iterator_types = ["parallel", "parallel", "parallel"]} ins(%0 : tensor<3x5x7xf32>) {
+    ^bb0(%arg2: f32):  // no predecessors
+      linalg.yield %arg2 : f32
+    } -> tensor<5x7x3xf32>
+    return %1 : tensor<5x7x3xf32>
+}
+
+// CHECK-LABEL: func @generic_op_120_permultation_reshape_producer_fusion
+//   CHECK-NOT: linalg.tensor_reshape
+//       CHECK: linalg.generic
+//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
+//   CHECK-NOT: linalg.tensor_reshape
+
+// -----
+
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d1, d0 * 7 + d2)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+
+#map0 = affine_map<(d0, d1, d2) -> (d0)>
+#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map2 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
+#map3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+func @generic_op_102_permultation_reshape_producer_fusion(%arg0 : tensor<3x35xf32>) -> tensor<5x3x7xf32> {
+  %0 = linalg.tensor_reshape %arg0 [#map0, #map1] : tensor<3x35xf32> into tensor<3x5x7xf32>
+  %1 = linalg.generic {indexing_maps = [#map2, #map3], iterator_types = ["parallel", "parallel", "parallel"]} ins(%0 : tensor<3x5x7xf32>) {
+    ^bb0(%arg2: f32):  // no predecessors
+      linalg.yield %arg2 : f32
+    } -> tensor<5x3x7xf32>
+    return %1 : tensor<5x3x7xf32>
+}
+
+// CHECK-LABEL: func @generic_op_102_permultation_reshape_producer_fusion
+//   CHECK-NOT: linalg.tensor_reshape
+//       CHECK: linalg.generic
+//  CHECK-SAME:   indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
+//   CHECK-NOT: linalg.tensor_reshape