From: Diego Caballero Date: Wed, 7 Aug 2019 17:31:14 +0000 (-0700) Subject: Fix verification of zero-dim memref in affine.load/affine.store/std.load/std.store X-Git-Url: http://review.tizen.org/git/?a=commitdiff_plain;h=c6a006d4c7987accbbed37b1462f2c7a4411e4bc;p=platform%2Fupstream%2Fllvm.git 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 --- diff --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h index 6a60223..711cfd8 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 9f5f873..3e4815a 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 767c2e3..9f347f9 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 f31868b..2ade7b9 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 f459c81..f2f4b2c 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 29df21e..ae55422 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 cfcf997..a2ea19e 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 9d90e33..da8b091 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 +}