[mlir][linalg] fix crash in vectorization of elementwise operations

[christopherbate]

Summary

The current vectorization logic implicitly expects "elementwise"
linalg ops to have projected permutations for indexing maps, but
the precondition logic misses this check. This can result in a
crash when executing the generic vectorization transform on an op
with a non-projected permutation input indexing map. This change
fixes the logic and adds a test (which crashes without this fix).

[christopherbate]

Fix typo

[ThomasRaoux]

Good catch!

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

if (!(isa<arith::ConstantOp, func::ConstantOp, linalg::YieldOp,
           OpTrait::hasElementwiseMappableTraits(&op)) ||
         llvm::any_of(op.getResultTypes(),
                      [](Type type) { return !type.isIntOrIndexOrFloat(); }))
       return false;
   }
   return true;
 }
 
+/// Returns `true` if all indexing maps of the linalg op are projected
+/// permutations.
+static bool allIndexingsAreProjectedPermutation(LinalgOp op) {
+  return llvm::all_of(op.getIndexingMaps(), [](AffineMap m) {
+    return m.isProjectedPermutation(/*allowZeroInResults=*/true);
+  });
+}
+
 // Return true if the op is an element-wise linalg op.
 static bool isElementwise(Operation *op) {
   auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
   if (!linalgOp)
     return false;
   if (linalgOp.getNumLoops() != linalgOp.getNumParallelLoops())
     return false;
+
+  if (!allIndexingsAreProjectedPermutation(linalgOp))
+    return false;
+
   // TODO: relax the restrictions on indexing map.
   for (OpOperand *opOperand : linalgOp.getOutputOperands()) {
     if (!linalgOp.getTiedIndexingMap(opOperand).isPermutation())
       return false;
   }
   return hasOnlyScalarElementwiseOp(linalgOp->getRegion(0));
 }
 
if (result.status == VectorizationStatus::NewOp) {
       LDBG("new vector op: " << *result.newOp;);
       bvm.map(op.getResults(), result.newOp->getResults());
     }
   }
 
   return success();
 }
 
-/// Helper function to vectorize a `linalgOp` with contraction semantics in a
-/// generic fashion.
-/// This helper is needed atm because the truly generic implementation requires
-/// good vector.multi_reduce folding patterns that are currently NYI.
-// TODO: drop reliance on a specific pattern.
-static bool allIndexingsAreProjectedPermutation(LinalgOp op) {
-  return llvm::all_of(op.getIndexingMaps(), [](AffineMap m) {
-    return m.isProjectedPermutation(/*allowZeroInResults=*/true);
-  });
-}
-
 // TODO: probably need some extra checks for reduction followed by consumer
 // ops that may not commute (e.g. linear reduction + non-linear instructions).
 static LogicalResult reductionPreconditions(LinalgOp op) {
   if (llvm::none_of(op.iterator_types(), isReductionIterator)) {
     LDBG("reduction precondition failed: no reduction iterator");
     return failure();
   }
   for (OpOperand *opOperand : op.getOutputOperands()) {

mlir/test/Dialect/Linalg/vectorization.mlir

%2 = linalg.generic {
          outs(%1 : tensor<f32>) {
     ^bb0(%a: f32, %b: f32):
       %3 = arith.addf %a, %b : f32
       linalg.yield %3 : f32
     } -> tensor<f32>
 
   return %2 : tensor<f32>
 }
+
+
+// -----
+
+// This test checks that vectorization does not occur when an input indexing map
+// is not a projected permutation. In the future, this can be converted to a 
+// positive test when support is added.
+
+// CHECK-LABEL:   func @not_projected_permutation
+func.func @not_projected_permutation(%arg0: tensor<8x8xf32>) -> tensor<6x6x3x3xf32> {
+  %c0 = arith.constant 0.0 : f32
+  %init = linalg.init_tensor [6, 6, 3, 3] : tensor<6x6x3x3xf32>
+  %fill = linalg.fill ins(%c0 : f32) outs(%init : tensor<6x6x3x3xf32>) -> tensor<6x6x3x3xf32>
+  // CHECK: linalg.generic
+  %result = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0 + d2, d1 + d3)>,
+                                             affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
+   iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+   ins(%arg0 : tensor<8x8xf32>)
+   outs(%fill : tensor<6x6x3x3xf32>) {
+    ^bb0(%arg7: f32, %arg9: f32):
+      linalg.yield %arg7 : f32
+    } -> tensor<6x6x3x3xf32>
+  return %result : tensor<6x6x3x3xf32>
+}