[mlir][Linalg] Disable const -> linalg.generic when fused op is illegal.

[mravishankar]

Summary

Fusing a constant with a linalg.generic operation can result in the
fused operation being illegal since the loop bound computation
fails. Avoid such fusions.

[asaadaldien]

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

667

Shouldn't we consider folding that into a splat constant tensor<3xf32>?

[mravishankar]

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

667

Yes, that would be a good idea. Doing this at Linalg level would be interesting, but seems harder. Seems easier to do at a higher level (say MHLO or TOSA).

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

for (auto operand : llvm::enumerate(linalgOp.getInputOpOperands())) {
 
       // The indexing_maps for the operands of the fused operation are same as
       // those for the operands of the linalgOp without the indexing map at
       // operand.index()
       SmallVector<AffineMap, 4> fusedIndexMaps = llvm::to_vector<4>(
           linalgOp.indexing_maps().getAsValueRange<AffineMapAttr>());
       fusedIndexMaps.erase(std::next(fusedIndexMaps.begin(), operand.index()));
 
+      // Check if the operation shapes to loops map is computable.
+      if (!inversePermutation(concatAffineMaps(fusedIndexMaps))) {
+        return rewriter.notifyMatchFailure(
+            linalgOp, "fused op loop bound computation failed");
+      }
+
       // The operands list is same as the linalgOp with the argument for
       // constant index dropped.
       SmallVector<Value, 4> fusedOperands(linalgOp.getInputs());
       fusedOperands.erase(std::next(fusedOperands.begin(), operand.index()));
 
       // Create a constant scalar value from the splat constant.
       Value scalarConstant = rewriter.create<ConstantOp>(
           constantOp.getLoc(),

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

 // CHECK-COUNT-2: linalg.generic
 func @index_op(%arg0: tensor<1x8xindex>, %arg1: tensor<1x8xindex>) -> tensor<1x8xindex> {
   %0 = linalg.generic {
     indexing_maps = [affine_map<(i, j) -> (i, j)>],
     iterator_types = ["parallel", "parallel"]}
   outs(%arg0 : tensor<1x8xindex>) {
   ^bb0(%a: index):   // no predecessors
     %2 = linalg.index 1 : index
     linalg.yield %2 : index
   } -> tensor<1x8xindex>
   %1 = linalg.generic {
     indexing_maps = [affine_map<(i, j) -> (i, j)>, affine_map<(i, j) -> (i, j)>],
     iterator_types = ["parallel", "parallel"]}
   ins(%0 : tensor<1x8xindex>)
   outs(%arg1 : tensor<1x8xindex>) {
   ^bb0(%a: index, %b: index):   // no predecessors
     %2 = linalg.index 0 : index
     %3 = addi %2, %a : index
     linalg.yield %3 : index
   } -> tensor<1x8xindex>
   return %1 : tensor<1x8xindex>
 }
+
+// -----
+
+// CHECK-LABEL: func @no_fuse_constant_with_reduction
+func @no_fuse_constant_with_reduction() -> tensor<3xf32>
+{
+  //      CHECK: %[[CONST:.+]] = constant {{.+}} : tensor<3x2xf32>
+  //      CHECK: %[[RESULT:.+]] = linalg.generic
+  // CHECK-SAME:   ins(%[[CONST]] : tensor<3x2xf32>)
+  //      CHECK: return %[[RESULT]]
+  %three = constant dense<3.0> : tensor<3x2xf32>
+  %init = linalg.init_tensor [3] : tensor<3xf32>
+  %result = linalg.generic {
+      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                       affine_map<(d0, d1) -> (d0)>],
+      iterator_types = ["parallel", "reduction"]}
+     ins(%three : tensor<3x2xf32>) outs(%init : tensor<3xf32>) {
+     ^bb0(%arg0 : f32, %arg1 : f32):
+        %0 = addf %arg0, %arg1 : f32
+        linalg.yield %0 : f32
+  } -> tensor<3xf32>
+  return %result : tensor<3xf32>
+}