Fix verification of zero-dim memref in affine.load/affine.store/std.load/std.store

Verification complained when using zero-dimensional memrefs in
affine.load, affine.store, std.load and std.store. This PR extends
verification so that those memrefs can be used.

Closes tensorflow/mlir#58

COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/mlir/pull/58 from dcaballe:dcaballe/zero-dim 49bcdcd45c52c48beca776431328e5ce551dfa9e
PiperOrigin-RevId: 262164916

mlir/include/mlir/IR/AffineMap.h

@@ -52,6 +52,9 @@ public:
   AffineMap(const AffineMap &other) : map(other.map) {}
   AffineMap &operator=(const AffineMap &other) = default;
 
+  /// Returns a zero result affine map with no dimensions or symbols: () -> ().
+  static AffineMap get(MLIRContext *context);
+
   static AffineMap get(unsigned dimCount, unsigned symbolCount,
                        ArrayRef<AffineExpr> results);
 
@@ -141,6 +144,9 @@ public:
 
 private:
   ImplType *map;
+
+  static AffineMap getImpl(unsigned dimCount, unsigned symbolCount,
+                           ArrayRef<AffineExpr> results, MLIRContext *context);
 };
 
 // Make AffineExpr hashable.

mlir/include/mlir/IR/Builders.h

@@ -150,6 +150,8 @@ public:
                          ArrayRef<AffineExpr> results);
 
   // Special cases of affine maps and integer sets
+  /// Returns a zero result affine map with no dimensions or symbols: () -> ().
+  AffineMap getEmptyAffineMap();
   /// Returns a single constant result affine map with 0 dimensions and 0
   /// symbols.  One constant result: () -> (val).
   AffineMap getConstantAffineMap(int64_t val);

mlir/lib/AffineOps/AffineOps.cpp

@@ -1591,7 +1591,11 @@ void AffineLoadOp::build(Builder *builder, OperationState *result,
   result->addOperands(memref);
   result->addOperands(indices);
   auto memrefType = memref->getType().cast<MemRefType>();
-  auto map = builder->getMultiDimIdentityMap(memrefType.getRank());
+  auto rank = memrefType.getRank();
+  // Create identity map for memrefs with at least one dimension or () -> ()
+  // for zero-dimensional memrefs.
+  auto map = rank ? builder->getMultiDimIdentityMap(rank)
+                  : builder->getEmptyAffineMap();
   result->addAttribute(getMapAttrName(), builder->getAffineMapAttr(map));
   result->types.push_back(memrefType.getElementType());
 }

mlir/lib/IR/Builders.cpp

@@ -297,6 +297,8 @@ IntegerSet Builder::getIntegerSet(unsigned dimCount, unsigned symbolCount,
   return IntegerSet::get(dimCount, symbolCount, constraints, isEq);
 }
 
+AffineMap Builder::getEmptyAffineMap() { return AffineMap::get(context); }
+
 AffineMap Builder::getConstantAffineMap(int64_t val) {
   return AffineMap::get(/*dimCount=*/0, /*symbolCount=*/0,
                         {getAffineConstantExpr(val)});

mlir/lib/IR/MLIRContext.cpp

@@ -568,12 +568,10 @@ StorageUniquer &MLIRContext::getAffineUniquer() {
   return getImpl().affineUniquer;
 }
 
-AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,
-                         ArrayRef<AffineExpr> results) {
-  // The number of results can't be zero.
-  assert(!results.empty());
-
-  auto &impl = results[0].getContext()->getImpl();
+AffineMap AffineMap::getImpl(unsigned dimCount, unsigned symbolCount,
+                             ArrayRef<AffineExpr> results,
+                             MLIRContext *context) {
+  auto &impl = context->getImpl();
   auto key = std::make_tuple(dimCount, symbolCount, results);
 
   // Safely get or create an AffineMap instance.
@@ -589,6 +587,17 @@ AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,
   });
 }
 
+AffineMap AffineMap::get(MLIRContext *context) {
+  return getImpl(/*dimCount=*/0, /*symbolCount=*/0, /*results=*/{}, context);
+}
+
+AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,
+                         ArrayRef<AffineExpr> results) {
+  // The number of results can't be zero.
+  assert(!results.empty());
+  return getImpl(dimCount, symbolCount, results, results[0].getContext());
+}
+
 //===----------------------------------------------------------------------===//
 // Integer Sets: these are allocated into the bump pointer, and are immutable.
 // Unlike AffineMap's, these are uniqued only if they are small.

mlir/test/AffineOps/load-store.mlir

@@ -182,4 +182,15 @@ func @test7() {
 // CHECK: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   }
   return
-}
+}
+
+// -----
+
+// Test with zero-dimensional operands.
+func @zero_dim(%arg0 : memref<i32>, %arg1 : memref<i32>) {
+  %0 = affine.load %arg0[] : memref<i32>
+  affine.store %0, %arg1[] : memref<i32>
+  // CHECK: affine.load %{{.*}}[] : memref<i32>
+  // CHECK: affine.store %{{.*}}, %{{.*}}[] : memref<i32>
+  return
+}

mlir/test/IR/core-ops.mlir

@@ -339,6 +339,16 @@ func @load_store(memref<4x4xi32>, index) {
   return
 }
 
+// Test with zero-dimensional operands using no index in load/store.
+// CHECK-LABEL: func @zero_dim_no_idx
+func @zero_dim_no_idx(%arg0 : memref<i32>, %arg1 : memref<i32>, %arg2 : memref<i32>) {
+  %0 = std.load %arg0[] : memref<i32>
+  std.store %0, %arg1[] : memref<i32>
+  return
+  // CHECK: %0 = load %{{.*}}[] : memref<i32>
+  // CHECK: store %{{.*}}, %{{.*}}[] : memref<i32>
+}
+
 // CHECK-LABEL: func @return_op(%arg0: i32) -> i32 {
 func @return_op(%a : i32) -> i32 {
   // CHECK: return %arg0 : i32

mlir/test/IR/invalid.mlir

@@ -1138,4 +1138,4 @@ func @integer_too_wide_in_tensor() {
 func @bool_literal_in_non_bool_tensor() {
   // expected-error @+1 {{expected i1 type for 'true' or 'false' values}}
   "foo"() {bar = dense<true> : tensor<2xi16>} : () -> ()
-}
+}