Add check to ensure pass fails gracefully

Author: Prakhar-Dixit

mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp

@@ -1217,6 +1217,21 @@ static Operation *vectorizeAffineLoad(AffineLoadOp loadOp,
     indices.append(mapOperands.begin(), mapOperands.end());
   }
 
+  for (auto &kvp : state.vecLoopToVecDim) {
+    AffineForOp forOp = cast<AffineForOp>(kvp.first);
+    auto invariants =
+        affine::getInvariantAccesses(forOp.getInductionVar(), indices);
+    unsigned nonInvariant = 0;
+    for (Value idx : indices)
+      if (!invariants.count(idx))
+        ++nonInvariant;
+    if (nonInvariant > 1) {
+      LLVM_DEBUG(dbgs() << "\n[early-vect] Bail out: loop IV "
+                        << forOp.getInductionVar() << " drives " << nonInvariant
+                        << " indices (must be ≤1)\n");
+      return nullptr;
+    }
+  }
   // Compute permutation map using the information of new vector loops.
   auto permutationMap = makePermutationMap(state.builder.getInsertionBlock(),
                                            indices, state.vecLoopToVecDim);
@@ -1262,6 +1277,21 @@ static Operation *vectorizeAffineStore(AffineStoreOp storeOp,
   else
     indices.append(mapOperands.begin(), mapOperands.end());
 
+  for (auto &kvp : state.vecLoopToVecDim) {
+    AffineForOp forOp = cast<AffineForOp>(kvp.first);
+    auto invariants =
+        affine::getInvariantAccesses(forOp.getInductionVar(), indices);
+    unsigned nonInvariant = 0;
+    for (Value idx : indices)
+      if (!invariants.count(idx))
+        ++nonInvariant;
+    if (nonInvariant > 1) {
+      LLVM_DEBUG(dbgs() << "\n[early-vect] Bail out: loop IV "
+                        << forOp.getInductionVar() << " drives " << nonInvariant
+                        << " indices (must be ≤1)\n");
+      return nullptr;
+    }
+  }
   // Compute permutation map using the information of new vector loops.
   auto permutationMap = makePermutationMap(state.builder.getInsertionBlock(),
                                            indices, state.vecLoopToVecDim);

mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp

@@ -141,14 +141,8 @@ static AffineMap makePermutationMap(
     unsigned countInvariantIndices = 0;
     for (unsigned dim = 0; dim < numIndices; ++dim) {
       if (!invariants.count(indices[dim])) {
-        if (perm[kvp.second] != getAffineConstantExpr(0, context)) {
-          auto loopOp = cast<affine::AffineForOp>(kvp.first);
-          loopOp->emitError(
-              "loop induction variable is used in multiple indices, which is "
-              "unsupported for vectorization. Consider using nested loops "
-              "instead of a single loop with affine.apply.");
-          return AffineMap();
-        }
+        assert(perm[kvp.second] == getAffineConstantExpr(0, context) &&
+               "permutationMap already has an entry along dim");
         perm[kvp.second] = getAffineDimExpr(dim, context);
       } else {
         ++countInvariantIndices;

mlir/test/Dialect/Affine/SuperVectorize/vectorize_unsupported.mlir

@@ -15,7 +15,7 @@ func.func @unparallel_loop_reduction_unsupported(%in: memref<256x512xf32>, %out:
 #map = affine_map<(d0)[s0] -> (d0 mod s0)>
 #map1 = affine_map<(d0)[s0] -> (d0 floordiv s0)>
 
-func.func @single_loop_unrolling_2D_access_pattern(%arg0: index) -> memref<2x2xf32> {
+func.func @single_loop_unrolling_2D_access_pattern_storeOp(%arg0: index) -> memref<2x2xf32> {
   %c2 = arith.constant 2 : index
   %cst = arith.constant 1.0 : f32
   %alloc = memref.alloc() : memref<2x2xf32>
@@ -33,7 +33,7 @@ func.func @single_loop_unrolling_2D_access_pattern(%arg0: index) -> memref<2x2xf
 // CHECK: #[[$ATTR_1:.+]] = affine_map<(d0)[s0] -> (d0 mod s0)>
 
 // CHECK-LABEL:   func.func @single_loop_unrolling_2D_access_pattern(
-// CHECK-SAME:                            %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index) -> memref<2x2xf32> {
+// CHECK-SAME:      %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index) -> memref<2x2xf32> {
 // CHECK:           %[[VAL_1:.*]] = arith.constant 2 : index
 // CHECK:           %[[VAL_2:.*]] = arith.constant 1.000000e+00 : f32
 // CHECK:           %[[VAL_3:.*]] = memref.alloc() : memref<2x2xf32>
@@ -43,4 +43,37 @@ func.func @single_loop_unrolling_2D_access_pattern(%arg0: index) -> memref<2x2xf
 // CHECK:             affine.store %[[VAL_2]], %[[VAL_3]]{{\[}}%[[VAL_5]], %[[VAL_6]]] : memref<2x2xf32>
 // CHECK:           }
 // CHECK:           return %[[VAL_3]] : memref<2x2xf32>
+// CHECK:         }
+
+// -----
+
+#map = affine_map<(d0)[s0] -> (d0 mod s0)>
+#map1 = affine_map<(d0)[s0] -> (d0 floordiv s0)>
+
+func.func @single_loop_unrolling_2D_access_pattern_loadOp(%arg0: index) -> memref<2x2xf32> {
+  %c2 = arith.constant 2 : index
+  %alloc = memref.alloc() : memref<2x2xf32>
+
+  affine.for %i = 0 to 4 {
+    %row = affine.apply #map1(%i)[%c2]  
+    %col = affine.apply #map(%i)[%c2]  
+    %val = affine.load %alloc[%row, %col] : memref<2x2xf32>
+  }
+
+  return %alloc : memref<2x2xf32>
+}
+
+// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0)[s0] -> (d0 floordiv s0)>
+// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0)[s0] -> (d0 mod s0)>
+
+// CHECK-LABEL:   func.func @single_loop_unrolling_2D_access_pattern(
+// CHECK-SAME:      %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: index) -> memref<2x2xf32> {
+// CHECK:           %[[VAL_1:.*]] = arith.constant 2 : index
+// CHECK:           %[[VAL_2:.*]] = memref.alloc() : memref<2x2xf32>
+// CHECK:           affine.for %[[VAL_3:.*]] = 0 to 4 {
+// CHECK:             %[[VAL_4:.*]] = affine.apply #[[$ATTR_0]](%[[VAL_3]]){{\[}}%[[VAL_1]]]
+// CHECK:             %[[VAL_5:.*]] = affine.apply #[[$ATTR_1]](%[[VAL_3]]){{\[}}%[[VAL_1]]]
+// CHECK:             %[[VAL_6:.*]] = affine.load %[[VAL_2]]{{\[}}%[[VAL_4]], %[[VAL_5]]] : memref<2x2xf32>
+// CHECK:           }
+// CHECK:           return %[[VAL_2]] : memref<2x2xf32>
 // CHECK:         }