[mlir][Affine] Fix vectorizability check for multiple load/stores

[sgrechanik]

Summary

This patch fixes a bug that allowed vectorizing of loops with loads and
stores having indexing functions varying along different memory
dimensions.

[aartbik]

mlir/lib/Analysis/LoopAnalysis.cpp

344

note that we have a subtle difference of not setting memRefDim for -1 after this change;
that is probably harmless, but perhaps you want this out of the if to preserve behavior?

mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir

402

can you please indent the CHECKs that are inside the method body?

[sgrechanik]

• Set *memRefDim to -1 in the beginning
• Fix indentation for the test case

mlir/lib/Analysis/LoopAnalysis.cpp

344

I think that if we have some memory ops that are invariant along the induction variable and some ops that vary along the same memory dimension N, then we should write N to *memRefDim. So we cannot write -1 to *memRefDim if isContiguousAccess returned -1, otherwise we may overwrite N.
However we should probably set *memRefDim to -1 at the beginning of this function in case the caller didn't do this.

[dcaballe]

Thanks for fixing this, Sergei! LGTM. I'll leave the final approval to Aart or Nicolas.

mlir/lib/Analysis/LoopAnalysis.cpp

if (isVectorizableOp && !isVectorizableOp(loop, *op)) {
       return false;
     }
   }
   return true;
 }
 
 bool mlir::isVectorizableLoopBody(AffineForOp loop, int *memRefDim,
                                   NestedPattern &vectorTransferMatcher) {
+  *memRefDim = -1;
   VectorizableOpFun fun([memRefDim](AffineForOp loop, Operation &op) {
     auto load = dyn_cast<AffineLoadOp>(op);
     auto store = dyn_cast<AffineStoreOp>(op);
-    return load ? isContiguousAccess(loop.getInductionVar(), load, memRefDim)
-                : isContiguousAccess(loop.getInductionVar(), store, memRefDim);
+    int thisOpMemRefDim = -1;
+    bool isContiguous = load ? isContiguousAccess(loop.getInductionVar(), load,
+                                                  &thisOpMemRefDim)
+                             : isContiguousAccess(loop.getInductionVar(), store,
+                                                  &thisOpMemRefDim);
+    if (thisOpMemRefDim != -1) {
+      // If memory accesses vary across different dimensions then the loop is
+      // not vectorizable.
+      if (*memRefDim != -1 && *memRefDim != thisOpMemRefDim)
+        return false;
+      *memRefDim = thisOpMemRefDim;
+    }
+    return isContiguous;
   });
   return isVectorizableLoopBodyWithOpCond(loop, fun, vectorTransferMatcher);
 }
 
 bool mlir::isVectorizableLoopBody(AffineForOp loop,
                                   NestedPattern &vectorTransferMatcher) {
   return isVectorizableLoopBodyWithOpCond(loop, nullptr, vectorTransferMatcher);
 }

mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir

// CHECK:  affine.for %{{.*}}{{[0-9]*}} = 0 to %{{[0-9]*}} {
    affine.for %i15 = 0 to %M { // not vectorized due to condition below
      affine.if #set0(%i15) {
        %a15 = affine.load %A[%c0, %c0] : memref<?x?xf32>
      }
    }
    return
 }
 
+// CHECK-LABEL: func @vec_rejected_11
+func @vec_rejected_11(%A : memref<?x?xf32>, %B : memref<?x?x?xf32>) {
+  // CHECK-DAG: %[[C0:.*]] = constant 0 : index
+  // CHECK-DAG: %[[C1:.*]] = constant 1 : index
+  // CHECK-DAG: %[[C2:.*]] = constant 2 : index
+  // CHECK-DAG: [[ARG_M:%[0-9]+]] = dim %{{.*}}, %[[C0]] : memref<?x?xf32>
+  // CHECK-DAG: [[ARG_N:%[0-9]+]] = dim %{{.*}}, %[[C1]] : memref<?x?xf32>
+  // CHECK-DAG: [[ARG_P:%[0-9]+]] = dim %{{.*}}, %[[C2]] : memref<?x?x?xf32>
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %M = dim %A, %c0 : memref<?x?xf32>
+  %N = dim %A, %c1 : memref<?x?xf32>
+  %P = dim %B, %c2 : memref<?x?x?xf32>
+
+  // CHECK: for [[IV10:%[arg0-9]*]] = 0 to %{{[0-9]*}} {
+  // CHECK:   for [[IV11:%[arg0-9]*]] = 0 to %{{[0-9]*}} {
+  // This is similar to vec_rejected_5, but the order of indices is different.
+  affine.for %i10 = 0 to %M { // not vectorized
+    affine.for %i11 = 0 to %N { // not vectorized
+      %a11 = affine.load %A[%i11, %i10] : memref<?x?xf32>
+      affine.store %a11, %A[%i10, %i11] : memref<?x?xf32>
+    }
+  }
+  return
+}
+
 // This should not vectorize due to the sequential dependence in the scf.
 // CHECK-LABEL: @vec_rejected_sequential
 func @vec_rejected_sequential(%A : memref<?xf32>) {
   %c0 = constant 0 : index
   %N = dim %A, %c0 : memref<?xf32>
   affine.for %i = 0 to %N {
     // CHECK-NOT: vector
     %a = affine.load %A[%i] : memref<?xf32>
     // CHECK-NOT: vector
     affine.store %a, %A[%i + 1] : memref<?xf32>
   }
   return
 }