From 9e6cf0d0258aa79b759db9c68a447071ab6ee4ab Mon Sep 17 00:00:00 2001 From: Diego Caballero <diego.caballero@intel.com> Date: Tue, 20 Aug 2019 10:43:45 -0700 Subject: [PATCH] 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 | 3 +++ mlir/lib/AffineOps/AffineOps.cpp | 6 +++++- mlir/lib/IR/AffineMap.cpp | 6 +++++- mlir/lib/IR/AffineMapDetail.h | 2 ++ mlir/lib/IR/AsmPrinter.cpp | 2 -- mlir/lib/IR/Attributes.cpp | 3 +-- mlir/lib/IR/MLIRContext.cpp | 2 +- mlir/test/EDSC/builder-api-test.cpp | 32 +++++++++++++++++++++++++++++ 8 files changed, 49 insertions(+), 7 deletions(-) diff --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h index 711cfd889807..b1ab50f937a3 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 f3af9599b595..b00a11083ecf 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 1b6bbe57c90a..e56d0e83f652 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 af1d89cd2395..a247783540c9 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 6825cb8ad5fa..e01944b1f7b1 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 507417a9f903..a8101a289908 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 f2f4b2c9a4e2..0551be59edd5 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 386d744ef329..f4c04eeb7fb7 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;