Fix build of affine load/store with empty map

tensorflow/mlir#58 fixed and exercised
verification of load/store ops using empty affine maps. Unfortunately,
it didn't exercise the creation of them. This PR addresses that aspect.
It removes the assumption of AffineMap having at least one result and
stores a pointer to MLIRContext as member of AffineMap.

* Add empty map support to affine.store + test
* Move MLIRContext to AffineMapStorage

Closes tensorflow/mlir#74

PiperOrigin-RevId: 264416260

mlir/include/mlir/IR/AffineMap.h

@@ -76,6 +76,9 @@ public:
   /// dimensional identifiers.
   bool isIdentity() const;
 
+  /// Returns true if this affine map is an empty map, i.e., () -> ().
+  bool isEmpty() const;
+
   /// Returns true if this affine map is a single result constant function.
   bool isSingleConstant() const;
 

mlir/lib/AffineOps/AffineOps.cpp

@@ -1684,7 +1684,11 @@ void AffineStoreOp::build(Builder *builder, OperationState *result,
   result->addOperands(memref);
   result->addOperands(operands);
   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));
 }
 

mlir/lib/IR/AffineMap.cpp

@@ -115,7 +115,7 @@ AffineMap AffineMap::getMultiDimIdentityMap(unsigned numDims,
   return get(/*dimCount=*/numDims, /*symbolCount=*/0, dimExprs);
 }
 
-MLIRContext *AffineMap::getContext() const { return getResult(0).getContext(); }
+MLIRContext *AffineMap::getContext() const { return map->context; }
 
 bool AffineMap::isIdentity() const {
   if (getNumDims() != getNumResults())
@@ -129,6 +129,10 @@ bool AffineMap::isIdentity() const {
   return true;
 }
 
+bool AffineMap::isEmpty() const {
+  return getNumDims() == 0 && getNumSymbols() == 0 && getNumResults() == 0;
+}
+
 bool AffineMap::isSingleConstant() const {
   return getNumResults() == 1 && getResult(0).isa<AffineConstantExpr>();
 }

mlir/lib/IR/AffineMapDetail.h

@@ -36,6 +36,8 @@ struct AffineMapStorage {
   /// The affine expressions for this (multi-dimensional) map.
   /// TODO: use trailing objects for this.
   ArrayRef<AffineExpr> results;
+
+  MLIRContext *context;
 };
 
 } // end namespace detail

mlir/lib/IR/AsmPrinter.cpp

@@ -1096,8 +1096,6 @@ void ModulePrinter::printAffineMap(AffineMap map) {
     os << ']';
   }
 
-  // AffineMap should have at least one result.
-  assert(!map.getResults().empty());
   // Result affine expressions.
   os << " -> (";
   interleaveComma(map.getResults(),

mlir/lib/IR/Attributes.cpp

@@ -62,8 +62,7 @@ Dialect &Attribute::getDialect() const { return impl->getDialect(); }
 //===----------------------------------------------------------------------===//
 
 AffineMapAttr AffineMapAttr::get(AffineMap value) {
-  return Base::get(value.getResult(0).getContext(),
-                   StandardAttributes::AffineMap, value);
+  return Base::get(value.getContext(), StandardAttributes::AffineMap, value);
 }
 
 AffineMap AffineMapAttr::getValue() const { return getImpl()->value; }

mlir/lib/IR/MLIRContext.cpp

@@ -582,7 +582,7 @@ AffineMap AffineMap::getImpl(unsigned dimCount, unsigned symbolCount,
     results = copyArrayRefInto(impl.affineAllocator, results);
 
     // Initialize the memory using placement new.
-    new (res) detail::AffineMapStorage{dimCount, symbolCount, results};
+    new (res) detail::AffineMapStorage{dimCount, symbolCount, results, context};
     return AffineMap(res);
   });
 }

mlir/test/EDSC/builder-api-test.cpp

@@ -714,6 +714,38 @@ TEST_FUNC(indirect_access) {
   f.erase();
 }
 
+// Exercise affine loads and stores build with empty maps.
+TEST_FUNC(empty_map_load_store) {
+  using namespace edsc;
+  using namespace edsc::intrinsics;
+  using namespace edsc::op;
+  auto memrefType =
+      MemRefType::get({}, FloatType::getF32(&globalContext()), {}, 0);
+  auto f = makeFunction("empty_map_load_store", {},
+                        {memrefType, memrefType, memrefType, memrefType});
+
+  OpBuilder builder(f.getBody());
+  ScopedContext scope(builder, f.getLoc());
+  ValueHandle zero = constant_index(0);
+  ValueHandle one = constant_index(1);
+  IndexedValue input(f.getArgument(0)), res(f.getArgument(1));
+  IndexHandle iv;
+
+  // clang-format off
+  LoopBuilder(&iv, zero, one, 1)([&]{
+      res() = input();
+  });
+  // clang-format on
+
+  // clang-format off
+  // CHECK-LABEL: func @empty_map_load_store(
+  // CHECK:  [[A:%.*]] = affine.load %{{.*}}[]
+  // CHECK:  affine.store [[A]], %{{.*}}[]
+  // clang-format on
+  f.print(llvm::outs());
+  f.erase();
+}
+
 int main() {
   RUN_TESTS();
   return 0;