From c6a006d4c7987accbbed37b1462f2c7a4411e4bc Mon Sep 17 00:00:00 2001 From: Diego Caballero Date: Wed, 7 Aug 2019 10:31:14 -0700 Subject: [PATCH] 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 | 6 ++++++ mlir/include/mlir/IR/Builders.h | 2 ++ mlir/lib/AffineOps/AffineOps.cpp | 6 +++++- mlir/lib/IR/Builders.cpp | 2 ++ mlir/lib/IR/MLIRContext.cpp | 21 +++++++++++++++------ mlir/test/AffineOps/load-store.mlir | 13 ++++++++++++- mlir/test/IR/core-ops.mlir | 10 ++++++++++ mlir/test/IR/invalid.mlir | 2 +- 8 files changed, 53 insertions(+), 9 deletions(-) diff --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h index 6a6022365792..711cfd889807 100644 --- a/mlir/include/mlir/IR/AffineMap.h +++ b/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 results); @@ -141,6 +144,9 @@ public: private: ImplType *map; + + static AffineMap getImpl(unsigned dimCount, unsigned symbolCount, + ArrayRef results, MLIRContext *context); }; // Make AffineExpr hashable. diff --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h index 9f5f873d506b..3e4815a5f32b 100644 --- a/mlir/include/mlir/IR/Builders.h +++ b/mlir/include/mlir/IR/Builders.h @@ -150,6 +150,8 @@ public: ArrayRef 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); diff --git a/mlir/lib/AffineOps/AffineOps.cpp b/mlir/lib/AffineOps/AffineOps.cpp index 767c2e344d9a..9f347f9c15c8 100644 --- a/mlir/lib/AffineOps/AffineOps.cpp +++ b/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(); - 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()); } diff --git a/mlir/lib/IR/Builders.cpp b/mlir/lib/IR/Builders.cpp index f31868bd7946..2ade7b9f28a4 100644 --- a/mlir/lib/IR/Builders.cpp +++ b/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)}); diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp index f459c81f2331..f2f4b2c9a4e2 100644 --- a/mlir/lib/IR/MLIRContext.cpp +++ b/mlir/lib/IR/MLIRContext.cpp @@ -568,12 +568,10 @@ StorageUniquer &MLIRContext::getAffineUniquer() { return getImpl().affineUniquer; } -AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount, - ArrayRef 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 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 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. diff --git a/mlir/test/AffineOps/load-store.mlir b/mlir/test/AffineOps/load-store.mlir index 29df21e4aba1..ae554220589d 100644 --- a/mlir/test/AffineOps/load-store.mlir +++ b/mlir/test/AffineOps/load-store.mlir @@ -182,4 +182,15 @@ func @test7() { // CHECK: affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32> } return -} \ No newline at end of file +} + +// ----- + +// Test with zero-dimensional operands. +func @zero_dim(%arg0 : memref, %arg1 : memref) { + %0 = affine.load %arg0[] : memref + affine.store %0, %arg1[] : memref + // CHECK: affine.load %{{.*}}[] : memref + // CHECK: affine.store %{{.*}}, %{{.*}}[] : memref + return +} diff --git a/mlir/test/IR/core-ops.mlir b/mlir/test/IR/core-ops.mlir index cfcf99794903..a2ea19e6b263 100644 --- a/mlir/test/IR/core-ops.mlir +++ b/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, %arg1 : memref, %arg2 : memref) { + %0 = std.load %arg0[] : memref + std.store %0, %arg1[] : memref + return + // CHECK: %0 = load %{{.*}}[] : memref + // CHECK: store %{{.*}}, %{{.*}}[] : memref +} + // CHECK-LABEL: func @return_op(%arg0: i32) -> i32 { func @return_op(%a : i32) -> i32 { // CHECK: return %arg0 : i32 diff --git a/mlir/test/IR/invalid.mlir b/mlir/test/IR/invalid.mlir index 9d90e33fd352..da8b0911b593 100644 --- a/mlir/test/IR/invalid.mlir +++ b/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 : tensor<2xi16>} : () -> () -} \ No newline at end of file +}