8 files changed, 49 insertions, 7 deletions
diff --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h
index 711cfd88980..b1ab50f937a 100644
--- a/mlir/include/mlir/IR/AffineMap.h
+++ b/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;
 
diff --git a/mlir/lib/AffineOps/AffineOps.cpp b/mlir/lib/AffineOps/AffineOps.cpp
index f3af9599b59..b00a11083ec 100644
--- a/mlir/lib/AffineOps/AffineOps.cpp
+++ b/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));
 }
 
diff --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp
index 1b6bbe57c90..e56d0e83f65 100644
--- a/mlir/lib/IR/AffineMap.cpp
+++ b/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>();
 }
diff --git a/mlir/lib/IR/AffineMapDetail.h b/mlir/lib/IR/AffineMapDetail.h
index af1d89cd239..a247783540c 100644
--- a/mlir/lib/IR/AffineMapDetail.h
+++ b/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
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index 6825cb8ad5f..e01944b1f7b 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/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(),
diff --git a/mlir/lib/IR/Attributes.cpp b/mlir/lib/IR/Attributes.cpp
index 507417a9f90..a8101a28990 100644
--- a/mlir/lib/IR/Attributes.cpp
+++ b/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; }
diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index f2f4b2c9a4e..0551be59edd 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/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);
   });
 }
diff --git a/mlir/test/EDSC/builder-api-test.cpp b/mlir/test/EDSC/builder-api-test.cpp
index 386d744ef32..f4c04eeb7fb 100644
--- a/mlir/test/EDSC/builder-api-test.cpp
+++ b/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;