[mlir][Linalg] Disallow ops with index semantics in `PushExpandingReshape`.

[mravishankar]

Summary

This pattern is not written to handle operations with linalg.index
operations in its body, i.e. operations that have index semantics.

[ThomasRaoux]

Thanks Mahesh

[rriddle]

Looks like this is missing a test

[mravishankar]

In D117856#3261869, @rriddle wrote:

Looks like this is missing a test

I am not sure this needs a test. Its avoiding application of a pattern on ops that have indexing semantics. It is one of the preconditions for the pattern (like op having tensor semantics or all parallel iterators).

[rriddle]

In D117856#3261881, @mravishankar wrote:

In D117856#3261869, @rriddle wrote:

Looks like this is missing a test

I am not sure this needs a test. Its avoiding application of a pattern on ops that have indexing semantics. It is one of the preconditions for the pattern (like op having tensor semantics or all parallel iterators).

This is changing the behavior of the pattern though right? I would expect a new test that checks that this pattern isn't applied on those ops.

[mravishankar]

In D117856#3261884, @rriddle wrote:

In D117856#3261881, @mravishankar wrote:

In D117856#3261869, @rriddle wrote:

Looks like this is missing a test

I am not sure this needs a test. Its avoiding application of a pattern on ops that have indexing semantics. It is one of the preconditions for the pattern (like op having tensor semantics or all parallel iterators).

This is changing the behavior of the pattern though right? I would expect a new test that checks that this pattern isn't applied on those ops.

It is changing in the sense this was the oversight when initially written. I can add a test, but just for clarification, there are no tests for whether the pattern applies on generic op without tensor semantics or on generic op without all parallel loops. Its basically preconditions where the pattern should trigger. It was missed when the pattern was written initially. I am not trying to be difficult. I'll add a test if you suggest that it indeed needs a test. From my perspective it is not changing the pattern behavior in the sense that it wasn't meant to behave that way to start with.

[mehdi_amini]

I'll add a test if you suggest that it indeed needs a test. From my perspective it is not changing the pattern behavior in the sense that it wasn't meant to behave that way to start with.

But isn't what you're describing "fixing a bug"? I think River sees it this way and is asking for a "non regression test".
At least that's how I understand it, maybe you can elaborate against this?

[mravishankar]

Add test.

[mravishankar]

In D117856#3261992, @mehdi_amini wrote:

I'll add a test if you suggest that it indeed needs a test. From my perspective it is not changing the pattern behavior in the sense that it wasn't meant to behave that way to start with.

But isn't what you're describing "fixing a bug"? I think River sees it this way and is asking for a "non regression test".
At least that's how I understand it, maybe you can elaborate against this?

I added the test :) .

But to me this precondition check is at the same level as checking the op name before applying a pattern. This pattern is illegally triggering on an op it shouldnt have to start with cause the preconditions were wrong. In any case, there is no black and white answer here, so adding a test is cheaper :) .

[mravishankar]

Rebase.

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

 ///  %3 = tensor.expand_shape %2 [[0, 1], [2]]
 ///    : tensor<12544x16xf32> into tensor<112x112x16xf32>
 struct PushExpandingReshape : public OpRewritePattern<GenericOp> {
   using OpRewritePattern<GenericOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
     // Only apply to elementwise linalg on tensor.
-    if (!genericOp.hasTensorSemantics() ||
+    if (!genericOp.hasTensorSemantics() || genericOp.hasIndexSemantics() ||
         genericOp.getNumParallelLoops() != genericOp.getNumLoops())
       return failure();
     // Only support identity output maps. It could be extended to permuations if
     // needed.
     if (llvm::any_of(genericOp.getOutputOperands(), [&](OpOperand *opOperand) {
           return !genericOp.getTiedIndexingMap(opOperand).isIdentity();
         }))
       return failure();

mlir/test/Dialect/Linalg/fusion-push-reshape.mlir

func @type_correctness(%arg0 : tensor<6x5xi32>, %arg1 : tensor<5xf32>,
   return %28 : tensor<2x3x5xf32>
 }
 // CHECK-LABEL: func @type_correctness
 //       CHECK:   %[[OP:.+]] = linalg.generic
 //  CHECK-SAME:   ins(%{{.+}}, %{{.+}}, %{{.+}} : tensor<6x5xi32>, tensor<5xf32>, tensor<5xf32>)
 //  CHECK-SAME:   outs(%{{.+}} : tensor<6x5xf32>)
 //       CHECK:   tensor.expand_shape %[[OP]]
 //  CHECK-SAME:   tensor<6x5xf32> into tensor<2x3x5xf32>
+
+// -----
+
+func @generic_op_index_semantics(%A: tensor<?x16xi64>, %B: tensor<16xi64>, %init: tensor<?x112x16xi64>) -> tensor<?x112x16xi64> {
+  %0 = tensor.expand_shape %A [[0, 1], [2]]
+      : tensor<?x16xi64> into tensor<?x112x16xi64>
+  %2 = linalg.generic {indexing_maps = [
+    affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d2)>,
+    affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
+    iterator_types = ["parallel", "parallel", "parallel"]}
+  ins(%0, %B : tensor<?x112x16xi64>, tensor<16xi64>)
+  outs(%init : tensor<?x112x16xi64>) {
+  ^bb0(%arg1: i64, %arg2: i64, %arg3: i64):  // no predecessors
+    %index = linalg.index 0 : index
+    %1 = arith.index_cast %index : index to i64
+    %add = arith.addi %arg1, %1 : i64
+    %s = arith.subi %add, %arg2 : i64
+    linalg.yield %s : i64
+  } -> tensor<?x112x16xi64>
+  return %2 : tensor<?x112x16xi64>
+}
+//      CHECK: func @generic_op_index_semantics
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<?x16xi64>
+//      CHECK:   %[[RESHAPE:.+]] = tensor.expand_shape %[[ARG0]]
+//      CHECK:   %[[RESULT:.+]] = linalg.generic
+// CHECK-SAME:       ins(%[[RESHAPE]]
+//      CHECK:   return %[[RESULT]]