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

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();

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

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