[MLIR][Affine] Fix affine memory op verifiers to check valid dim/symbols properly

[bondhugula]

The checks only checked if the operands were either valid dims and
symbols without checking which ones should be valid dims and which
should be valid symbols. The necessary methods to check already existed.

With this fix, for some existing tests, it leads to a more accurate
error message. In other cases, invalid IRs are caught as shown in the test case added.

Related to: #120189

[llvmbot]

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-spirv

@llvm/pr-subscribers-mlir-affine

Author: Uday Bondhugula (bondhugula)

Changes

The checks only checked if the operands were either valid dims and
symbols without checking which ones should be valid dims and which
should be valid symbols. The necessary methods to check already existed.

With this fix, for some existing tests, it leads to a more accurate
error message. In other cases, invalid IRs are caught.

---

Full diff: https://github.com/llvm/llvm-project/pull/127375.diff

4 Files Affected:

• (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+7-11)
• (modified) mlir/test/Conversion/FuncToSPIRV/func-ops-to-spirv.mlir (+1-1)
• (modified) mlir/test/Dialect/Affine/invalid.mlir (+2-2)
• (modified) mlir/test/Dialect/Affine/load-store-invalid.mlir (+17-2)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 147f5dd7a24b6..66da16075017f 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -3044,8 +3044,9 @@ void AffineLoadOp::print(OpAsmPrinter &p) {
 
 /// Verify common indexing invariants of affine.load, affine.store,
 /// affine.vector_load and affine.vector_store.
+template <typename AffineMemOpTy>
 static LogicalResult
-verifyMemoryOpIndexing(Operation *op, AffineMapAttr mapAttr,
+verifyMemoryOpIndexing(AffineMemOpTy op, AffineMapAttr mapAttr,
                        Operation::operand_range mapOperands,
                        MemRefType memrefType, unsigned numIndexOperands) {
   AffineMap map = mapAttr.getValue();
@@ -3054,14 +3055,12 @@ verifyMemoryOpIndexing(Operation *op, AffineMapAttr mapAttr,
   if (map.getNumInputs() != numIndexOperands)
     return op->emitOpError("expects as many subscripts as affine map inputs");
 
-  Region *scope = getAffineScope(op);
   for (auto idx : mapOperands) {
     if (!idx.getType().isIndex())
       return op->emitOpError("index to load must have 'index' type");
-    if (!isValidAffineIndexOperand(idx, scope))
-      return op->emitOpError(
-          "index must be a valid dimension or symbol identifier");
   }
+  if (failed(verifyDimAndSymbolIdentifiers(op, mapOperands, map.getNumDims())))
+    return failure();
 
   return success();
 }
@@ -3072,8 +3071,7 @@ LogicalResult AffineLoadOp::verify() {
     return emitOpError("result type must match element type of memref");
 
   if (failed(verifyMemoryOpIndexing(
-          getOperation(),
-          (*this)->getAttrOfType<AffineMapAttr>(getMapAttrStrName()),
+          *this, (*this)->getAttrOfType<AffineMapAttr>(getMapAttrStrName()),
           getMapOperands(), memrefType,
           /*numIndexOperands=*/getNumOperands() - 1)))
     return failure();
@@ -3189,8 +3187,7 @@ LogicalResult AffineStoreOp::verify() {
         "value to store must have the same type as memref element type");
 
   if (failed(verifyMemoryOpIndexing(
-          getOperation(),
-          (*this)->getAttrOfType<AffineMapAttr>(getMapAttrStrName()),
+          *this, (*this)->getAttrOfType<AffineMapAttr>(getMapAttrStrName()),
           getMapOperands(), memrefType,
           /*numIndexOperands=*/getNumOperands() - 2)))
     return failure();
@@ -4411,8 +4408,7 @@ static LogicalResult verifyVectorMemoryOp(Operation *op, MemRefType memrefType,
 LogicalResult AffineVectorLoadOp::verify() {
   MemRefType memrefType = getMemRefType();
   if (failed(verifyMemoryOpIndexing(
-          getOperation(),
-          (*this)->getAttrOfType<AffineMapAttr>(getMapAttrStrName()),
+          *this, (*this)->getAttrOfType<AffineMapAttr>(getMapAttrStrName()),
           getMapOperands(), memrefType,
           /*numIndexOperands=*/getNumOperands() - 1)))
     return failure();
diff --git a/mlir/test/Conversion/FuncToSPIRV/func-ops-to-spirv.mlir b/mlir/test/Conversion/FuncToSPIRV/func-ops-to-spirv.mlir
index a09f1697fd724..69200cb02713d 100644
--- a/mlir/test/Conversion/FuncToSPIRV/func-ops-to-spirv.mlir
+++ b/mlir/test/Conversion/FuncToSPIRV/func-ops-to-spirv.mlir
@@ -55,7 +55,7 @@ func.func @dim_index_out_of_bounds() {
   %alloc_4 = memref.alloc() : memref<4xi64>
   %dim = memref.dim %alloc_4, %c6 : memref<4xi64>
   %alloca_100 = memref.alloca() : memref<100xi64>
-  // expected-error@+1 {{'affine.vector_load' op index must be a valid dimension or symbol identifier}}
+  // expected-error@+1 {{'affine.vector_load' op operand cannot be used as a dimension id}}
   %70 = affine.vector_load %alloca_100[%dim] : memref<100xi64>, vector<31xi64>
   return
 }
diff --git a/mlir/test/Dialect/Affine/invalid.mlir b/mlir/test/Dialect/Affine/invalid.mlir
index 44e484b9ba598..dc061cd4eb3d7 100644
--- a/mlir/test/Dialect/Affine/invalid.mlir
+++ b/mlir/test/Dialect/Affine/invalid.mlir
@@ -25,7 +25,7 @@ func.func @affine_load_invalid_dim(%M : memref<10xi32>) {
   "unknown"() ({
   ^bb0(%arg: index):
     affine.load %M[%arg] : memref<10xi32>
-    // expected-error@-1 {{index must be a valid dimension or symbol identifier}}
+    // expected-error@-1 {{op operand cannot be used as a dimension id}}
     cf.br ^bb1
   ^bb1:
     cf.br ^bb1
@@ -517,7 +517,7 @@ func.func @dynamic_dimension_index() {
     %idx = "unknown.test"() : () -> (index)
     %memref = "unknown.test"() : () -> memref<?x?xf32>
     %dim = memref.dim %memref, %idx : memref<?x?xf32>
-    // expected-error @below {{op index must be a valid dimension or symbol identifier}}
+    // expected-error@below {{op operand cannot be used as a dimension id}}
     affine.load %memref[%dim, %dim] : memref<?x?xf32>
     "unknown.terminator"() : () -> ()
   }) : () -> ()
diff --git a/mlir/test/Dialect/Affine/load-store-invalid.mlir b/mlir/test/Dialect/Affine/load-store-invalid.mlir
index 01d6b25dee695..5e2789321fba8 100644
--- a/mlir/test/Dialect/Affine/load-store-invalid.mlir
+++ b/mlir/test/Dialect/Affine/load-store-invalid.mlir
@@ -37,7 +37,7 @@ func.func @load_non_affine_index(%arg0 : index) {
   %0 = memref.alloc() : memref<10xf32>
   affine.for %i0 = 0 to 10 {
     %1 = arith.muli %i0, %arg0 : index
-    // expected-error@+1 {{op index must be a valid dimension or symbol identifier}}
+    // expected-error@+1 {{op operand cannot be used as a dimension id}}
     %v = affine.load %0[%1] : memref<10xf32>
   }
   return
@@ -50,7 +50,7 @@ func.func @store_non_affine_index(%arg0 : index) {
   %1 = arith.constant 11.0 : f32
   affine.for %i0 = 0 to 10 {
     %2 = arith.muli %i0, %arg0 : index
-    // expected-error@+1 {{op index must be a valid dimension or symbol identifier}}
+    // expected-error@+1 {{op operand cannot be used as a dimension id}}
     affine.store %1, %0[%2] : memref<10xf32>
   }
   return
@@ -140,3 +140,18 @@ func.func @dma_wait_non_affine_tag_index(%arg0 : index) {
   }
   return
 }
+
+// -----
+
+func.func @invalid_symbol() {
+  %alloc = memref.alloc() {alignment = 64 : i64} : memref<23x26xf32>
+  affine.for %arg1 = 0 to 1 {
+    affine.for %arg2 = 0 to 26 {
+      affine.for %arg3 = 0 to 23 {
+        affine.load %alloc[symbol(%arg1) * 23 + symbol(%arg3), %arg2] : memref<23x26xf32>
+        // expected-error@above {{op operand cannot be used as a symbol}}
+      }
+    }
+  }
+  return
+}

[bondhugula]

Obvious fix to the verifiers, merging.