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

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

Author: bondhugula

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

@@ -3048,8 +3048,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();
@@ -3058,14 +3059,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();
 }
@@ -3076,8 +3075,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();
@@ -3193,8 +3191,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();
@@ -4425,8 +4422,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();

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
 }

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
@@ -535,7 +535,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"() : () -> ()
   }) : () -> ()

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
+}