[mlir][Linalg] Avoid generating illegal operations during elementwise fusion.

[mravishankar]

Summary

In some cases, fusion can produce illegal operations if after fusion
the range of some of the loops cannot be computed from shapes of its
operands. Check for this case and abort the fusion if this happens.

[nicolasvasilache]

mlir/test/Dialect/Linalg/fusion-elementwise-ops.mlir

968

do we need such complex bodies?

[mravishankar]

Rebase and reduce test complexity.

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

fuseElementwiseOpsImpl(GenericOp producer, OpOperand *consumerOpOperand,
   auto fusedOp = rewriter.create<GenericOp>(
       consumer.getLoc(), consumer->getResultTypes(),
       /*inputs=*/fusedOperands,
       // TODO: handle outputs.
       consumerOutputs, rewriter.getAffineMapArrayAttr(fusedIndexMaps),
       consumer.iterator_types(),
       /*doc=*/nullptr,
       /*library_call=*/nullptr);
+  if (!fusedOp.getShapesToLoopsMap()) {
+    // Fused op has invalid indexing maps. Typically this means something is off
+    // in the input, but going ahead here would result in verification errors.
+    // So cleanup and abort.
+    rewriter.eraseOp(fusedOp);
+    return llvm::None;
+  }
 
   // Construct an AffineMap from consumer loops to producer loops.
   // consumer loop -> tensor index
   AffineMap consumerResultIndexMap =
       consumer.getTiedIndexingMap(consumerOpOperand);
   // tensor index -> producer loop
   AffineMap invProducerResultIndexMap =
       inversePermutation(producerResultIndexMap);

mlir/test/Dialect/Linalg/fusion-elementwise-ops.mlir

%8 = linalg.generic {
     iterator_types = ["parallel", "reduction"]
   } ins(%5 : tensor<?x?xf32>) outs(%7 : tensor<?xf32>) {
   ^bb0(%arg2: f32, %arg3: f32):  
     %9 = arith.maxf %arg2, %arg3 : f32
     linalg.yield %9 : f32
   } -> tensor<?xf32>
   return %8 : tensor<?xf32>
 }
+
+// -----
+
+func @illegal_fusion(%arg0 : tensor<5000xi64>, %arg1 : tensor<5000xi32>) -> tensor<5000xi32> {
+  %c1_i32 = arith.constant 1 : i32
+  %0 = linalg.generic {
+        indexing_maps = [affine_map<(d0) -> (d0)>],
+        iterator_types = ["parallel"]}
+        outs(%arg0 : tensor<5000xi64>) {
+        ^bb0(%arg3: i64):  // no predecessors
+          %22 = linalg.index 0 : index
+          %23 = arith.index_cast %22 : index to i64
+          linalg.yield %23 : i64
+        } -> tensor<5000xi64>
+  %1 = linalg.init_tensor [5000] : tensor<5000xi32>
+  %2 = linalg.generic {
+        indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1)>],
+        iterator_types = ["parallel", "parallel"]}
+        ins(%0 : tensor<5000xi64>) outs(%1 : tensor<5000xi32>) {
+        ^bb0(%arg3: i64, %arg5: i32):  // no predecessors
+          %22 = arith.index_cast %arg3 : i64 to index
+          %23 = tensor.extract %arg1[%22] : tensor<5000xi32>
+          linalg.yield %23 : i32
+        } -> tensor<5000xi32>
+  return %2 : tensor<5000xi32>
+}
+// CHECK-LABEL: func @illegal_fusion(
+//       CHECK:   %[[PRODUCER:.+]] = linalg.generic
+//       CHECK:    linalg.generic
+//   CHECK-SAME:       ins(%[[PRODUCER]]