From 666b7222191c9ecbfd7abacdebcb7ca9ef661f98 Mon Sep 17 00:00:00 2001 From: linuxlonelyeagle <2020382038@qq.com> Date: Mon, 20 Jan 2025 10:38:35 +0800 Subject: [PATCH] if the result of a Pure operation that whose operands are dimensional identifiers,then their results are dimensional identifiers. --- mlir/docs/Dialects/Affine.md | 4 +- mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 22 ++- mlir/test/Dialect/Affine/invalid.mlir | 12 -- .../Dialect/Affine/load-store-invalid.mlir | 92 --------- mlir/test/Dialect/Affine/ops.mlir | 174 ++++++++++++++++++ 5 files changed, 190 insertions(+), 114 deletions(-) diff --git a/mlir/docs/Dialects/Affine.md b/mlir/docs/Dialects/Affine.md index 0b6d7747e8a6f..7306e5fbb37b7 100644 --- a/mlir/docs/Dialects/Affine.md +++ b/mlir/docs/Dialects/Affine.md @@ -83,8 +83,8 @@ location of the SSA use. Dimensions may be bound not only to anything that a symbol is bound to, but also to induction variables of enclosing [`affine.for`](#affinefor-mliraffineforop) and [`affine.parallel`](#affineparallel-mliraffineparallelop) operations, and the result -of an [`affine.apply` operation](#affineapply-mliraffineapplyop) (which recursively -may use other dimensions and symbols). +of a `Pure` operation whose operands are valid dimensional identifiers. +(which recursively may use other dimensions and symbols). ### Affine Expressions diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp index 147f5dd7a24b6..acbf66fa79a74 100644 --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -274,7 +274,8 @@ Region *mlir::affine::getAffineScope(Operation *op) { // conditions: // *) It is valid as a symbol. // *) It is an induction variable. -// *) It is the result of affine apply operation with dimension id arguments. +// *) It is the result of a `Pure` operation whose operands with dimension id +// *) arguments. bool mlir::affine::isValidDim(Value value) { // The value must be an index type. if (!value.getType().isIndex()) @@ -294,7 +295,8 @@ bool mlir::affine::isValidDim(Value value) { // conditions: // *) It is valid as a symbol. // *) It is an induction variable. -// *) It is the result of an affine apply operation with dimension id operands. +// *) It is the result of a `Pure` operation whose operands with dimension id +// *) operands. bool mlir::affine::isValidDim(Value value, Region *region) { // The value must be an index type. if (!value.getType().isIndex()) @@ -304,20 +306,24 @@ bool mlir::affine::isValidDim(Value value, Region *region) { if (isValidSymbol(value, region)) return true; - auto *op = value.getDefiningOp(); - if (!op) { + auto *defOp = value.getDefiningOp(); + if (!defOp) { // This value has to be a block argument for an affine.for or an // affine.parallel. auto *parentOp = llvm::cast(value).getOwner()->getParentOp(); return isa(parentOp); } - // Affine apply operation is ok if all of its operands are ok. - if (auto applyOp = dyn_cast(op)) - return applyOp.isValidDim(region); + // `Pure` operation is ok if all of its operands are ok. + if (isPure(defOp) && llvm::all_of(defOp->getOperands(), [&](Value operand) { + return affine::isValidDim(operand, region); + })) { + return true; + } + // The dim op is okay if its operand memref/tensor is defined at the top // level. - if (auto dimOp = dyn_cast(op)) + if (auto dimOp = dyn_cast(defOp)) return isTopLevelValue(dimOp.getShapedValue()); return false; } diff --git a/mlir/test/Dialect/Affine/invalid.mlir b/mlir/test/Dialect/Affine/invalid.mlir index 44e484b9ba598..b3b2ec5552482 100644 --- a/mlir/test/Dialect/Affine/invalid.mlir +++ b/mlir/test/Dialect/Affine/invalid.mlir @@ -225,18 +225,6 @@ func.func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) { // ----- -func.func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) { - affine.for %x = 0 to 7 { - %y = arith.addi %x, %x : index - // expected-error@+1 {{operand cannot be used as a dimension id}} - affine.parallel (%i, %j) = (0, 0) to (%y, 100) step (10, 10) { - } - } - return -} - -// ----- - func.func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) { affine.for %x = 0 to 7 { %y = arith.addi %x, %x : index diff --git a/mlir/test/Dialect/Affine/load-store-invalid.mlir b/mlir/test/Dialect/Affine/load-store-invalid.mlir index 01d6b25dee695..d8eac141cc52a 100644 --- a/mlir/test/Dialect/Affine/load-store-invalid.mlir +++ b/mlir/test/Dialect/Affine/load-store-invalid.mlir @@ -33,31 +33,6 @@ func.func @store_too_few_subscripts_map(%arg0: memref, %arg1: index, %v // ----- -func.func @load_non_affine_index(%arg0 : index) { - %0 = memref.alloc() : memref<10xf32> - affine.for %i0 = 0 to 10 { - %1 = arith.muli %i0, %arg0 : index - // expected-error@+1 {{op index must be a valid dimension or symbol identifier}} - %v = affine.load %0[%1] : memref<10xf32> - } - return -} - -// ----- - -func.func @store_non_affine_index(%arg0 : index) { - %0 = memref.alloc() : memref<10xf32> - %1 = arith.constant 11.0 : f32 - affine.for %i0 = 0 to 10 { - %2 = arith.muli %i0, %arg0 : index - // expected-error@+1 {{op index must be a valid dimension or symbol identifier}} - affine.store %1, %0[%2] : memref<10xf32> - } - return -} - -// ----- - func.func @invalid_prefetch_rw(%i : index) { %0 = memref.alloc() : memref<10xf32> // expected-error@+1 {{rw specifier has to be 'read' or 'write'}} @@ -73,70 +48,3 @@ func.func @invalid_prefetch_cache_type(%i : index) { affine.prefetch %0[%i], read, locality<0>, false : memref<10xf32> return } - -// ----- - -func.func @dma_start_non_affine_src_index(%arg0 : index) { - %0 = memref.alloc() : memref<100xf32> - %1 = memref.alloc() : memref<100xf32, 2> - %2 = memref.alloc() : memref<1xi32, 4> - %c0 = arith.constant 0 : index - %c64 = arith.constant 64 : index - affine.for %i0 = 0 to 10 { - %3 = arith.muli %i0, %arg0 : index - // expected-error@+1 {{op src index must be a valid dimension or symbol identifier}} - affine.dma_start %0[%3], %1[%i0], %2[%c0], %c64 - : memref<100xf32>, memref<100xf32, 2>, memref<1xi32, 4> - } - return -} - -// ----- - -func.func @dma_start_non_affine_dst_index(%arg0 : index) { - %0 = memref.alloc() : memref<100xf32> - %1 = memref.alloc() : memref<100xf32, 2> - %2 = memref.alloc() : memref<1xi32, 4> - %c0 = arith.constant 0 : index - %c64 = arith.constant 64 : index - affine.for %i0 = 0 to 10 { - %3 = arith.muli %i0, %arg0 : index - // expected-error@+1 {{op dst index must be a valid dimension or symbol identifier}} - affine.dma_start %0[%i0], %1[%3], %2[%c0], %c64 - : memref<100xf32>, memref<100xf32, 2>, memref<1xi32, 4> - } - return -} - -// ----- - -func.func @dma_start_non_affine_tag_index(%arg0 : index) { - %0 = memref.alloc() : memref<100xf32> - %1 = memref.alloc() : memref<100xf32, 2> - %2 = memref.alloc() : memref<1xi32, 4> - %c0 = arith.constant 0 : index - %c64 = arith.constant 64 : index - affine.for %i0 = 0 to 10 { - %3 = arith.muli %i0, %arg0 : index - // expected-error@+1 {{op tag index must be a valid dimension or symbol identifier}} - affine.dma_start %0[%i0], %1[%arg0], %2[%3], %c64 - : memref<100xf32>, memref<100xf32, 2>, memref<1xi32, 4> - } - return -} - -// ----- - -func.func @dma_wait_non_affine_tag_index(%arg0 : index) { - %0 = memref.alloc() : memref<100xf32> - %1 = memref.alloc() : memref<100xf32, 2> - %2 = memref.alloc() : memref<1xi32, 4> - %c0 = arith.constant 0 : index - %c64 = arith.constant 64 : index - affine.for %i0 = 0 to 10 { - %3 = arith.muli %i0, %arg0 : index - // expected-error@+1 {{op index must be a valid dimension or symbol identifier}} - affine.dma_wait %2[%3], %c64 : memref<1xi32, 4> - } - return -} diff --git a/mlir/test/Dialect/Affine/ops.mlir b/mlir/test/Dialect/Affine/ops.mlir index e3721806989bb..74ba098ce2748 100644 --- a/mlir/test/Dialect/Affine/ops.mlir +++ b/mlir/test/Dialect/Affine/ops.mlir @@ -409,3 +409,177 @@ func.func @arith_add_vaild_symbol_lower_bound(%arg : index) { // CHECK: affine.for %[[VAL_3:.*]] = #[[$ATTR_0]](%[[VAL_2]]){{\[}}%[[VAL_0]]] to 7 { // CHECK: } // CHECK: } + +// ----- + +// CHECK-LABEL: func @affine_parallel + +func.func @affine_parallel(%arg0 : index, %arg1 : index, %arg2 : index) { + affine.for %x = 0 to 7 { + %y = arith.addi %x, %x : index + affine.parallel (%i, %j) = (0, 0) to (%y, 100) step (10, 10) { + } + } + return +} + +// CHECK-NEXT: affine.for +// CHECK-SAME: %[[VAL_0:.*]] = 0 to 7 { +// CHECK: %[[VAL_1:.*]] = arith.addi %[[VAL_0]], %[[VAL_0]] : index +// CHECK: affine.parallel (%{{.*}}, %{{.*}}) = (0, 0) to (%[[VAL_1]], 100) step (10, 10) { +// CHECK: } +// CHECK: } + +// ----- + +func.func @load_non_affine_index(%arg0 : index) { + %0 = memref.alloc() : memref<10xf32> + affine.for %i0 = 0 to 10 { + %1 = arith.muli %i0, %arg0 : index + %v = affine.load %0[%1] : memref<10xf32> + } + return +} + +// CHECK-LABEL: func @load_non_affine_index +// CHECK-SAME: %[[VAL_0:.*]]: index) { +// CHECK: %[[VAL_1:.*]] = memref.alloc() : memref<10xf32> +// CHECK: affine.for %[[VAL_2:.*]] = 0 to 10 { +// CHECK: %[[VAL_3:.*]] = arith.muli %[[VAL_2]], %[[VAL_0]] : index +// CHECK: %{{.*}} = affine.load %[[VAL_1]]{{\[}}%[[VAL_3]]] : memref<10xf32> +// CHECK: } + +// ----- + +func.func @store_non_affine_index(%arg0 : index) { + %0 = memref.alloc() : memref<10xf32> + %1 = arith.constant 11.0 : f32 + affine.for %i0 = 0 to 10 { + %2 = arith.muli %i0, %arg0 : index + affine.store %1, %0[%2] : memref<10xf32> + } + return +} + +// CHECK-LABEL: func @store_non_affine_index +// CHECK-SAME: %[[VAL_0:.*]]: index) { +// CHECK: %[[VAL_1:.*]] = memref.alloc() : memref<10xf32> +// CHECK: %[[VAL_2:.*]] = arith.constant 1.100000e+01 : f32 +// CHECK: affine.for %[[VAL_3:.*]] = 0 to 10 { +// CHECK: %[[VAL_4:.*]] = arith.muli %[[VAL_3]], %[[VAL_0]] : index +// CHECK: affine.store %[[VAL_2]], %[[VAL_1]]{{\[}}%[[VAL_4]]] : memref<10xf32> +// CHECK: } + +// ----- + +func.func @dma_start_non_affine_src_index(%arg0 : index) { + %0 = memref.alloc() : memref<100xf32> + %1 = memref.alloc() : memref<100xf32, 2> + %2 = memref.alloc() : memref<1xi32, 4> + %c0 = arith.constant 0 : index + %c64 = arith.constant 64 : index + affine.for %i0 = 0 to 10 { + %3 = arith.muli %i0, %arg0 : index + affine.dma_start %0[%3], %1[%i0], %2[%c0], %c64 + : memref<100xf32>, memref<100xf32, 2>, memref<1xi32, 4> + } + return +} + +// CHECK-LABEL: func @dma_start_non_affine_src_index +// CHECK-SAME: %[[VAL_0:.*]]: index) { +// CHECK: %[[VAL_1:.*]] = memref.alloc() : memref<100xf32> +// CHECK: %[[VAL_2:.*]] = memref.alloc() : memref<100xf32, 2> +// CHECK: %[[VAL_3:.*]] = memref.alloc() : memref<1xi32, 4> +// CHECK: %[[VAL_4:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_5:.*]] = arith.constant 64 : index +// CHECK: affine.for %[[VAL_6:.*]] = 0 to 10 { +// CHECK: %[[VAL_7:.*]] = arith.muli %[[VAL_6]], %[[VAL_0]] : index +// CHECK: affine.dma_start %[[VAL_1]]{{\[}}%[[VAL_7]]], %[[VAL_2]]{{\[}}%[[VAL_6]]], %[[VAL_3]]{{\[}}%[[VAL_4]]], %[[VAL_5]] +// CHECK-SAME: : memref<100xf32>, memref<100xf32, 2>, memref<1xi32, 4> +// CHECK: } + +// ----- + +func.func @dma_start_non_affine_dst_index(%arg0 : index) { + %0 = memref.alloc() : memref<100xf32> + %1 = memref.alloc() : memref<100xf32, 2> + %2 = memref.alloc() : memref<1xi32, 4> + %c0 = arith.constant 0 : index + %c64 = arith.constant 64 : index + affine.for %i0 = 0 to 10 { + %3 = arith.muli %i0, %arg0 : index + affine.dma_start %0[%i0], %1[%3], %2[%c0], %c64 + : memref<100xf32>, memref<100xf32, 2>, memref<1xi32, 4> + } + return +} + +// CHECK-LABEL: func @dma_start_non_affine_dst_index +// CHECK-SAME: %[[VAL_0:.*]]: index) { +// CHECK: %[[VAL_1:.*]] = memref.alloc() : memref<100xf32> +// CHECK: %[[VAL_2:.*]] = memref.alloc() : memref<100xf32, 2> +// CHECK: %[[VAL_3:.*]] = memref.alloc() : memref<1xi32, 4> +// CHECK: %[[VAL_4:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_5:.*]] = arith.constant 64 : index +// CHECK: affine.for %[[VAL_6:.*]] = 0 to 10 { +// CHECK: %[[VAL_7:.*]] = arith.muli %[[VAL_6]], %[[VAL_0]] : index +// CHECK: affine.dma_start %[[VAL_1]]{{\[}}%[[VAL_6]]], %[[VAL_2]]{{\[}}%[[VAL_7]]], %[[VAL_3]]{{\[}}%[[VAL_4]]], %[[VAL_5]] +// CHECK-SAME: : memref<100xf32>, memref<100xf32, 2>, memref<1xi32, 4> +// CHECK: } + +// ----- + +func.func @dma_start_non_affine_tag_index(%arg0 : index) { + %0 = memref.alloc() : memref<100xf32> + %1 = memref.alloc() : memref<100xf32, 2> + %2 = memref.alloc() : memref<1xi32, 4> + %c0 = arith.constant 0 : index + %c64 = arith.constant 64 : index + affine.for %i0 = 0 to 10 { + %3 = arith.muli %i0, %arg0 : index + affine.dma_start %0[%i0], %1[%arg0], %2[%3], %c64 + : memref<100xf32>, memref<100xf32, 2>, memref<1xi32, 4> + } + return +} + +// CHECK-LABEL: func @dma_start_non_affine_tag_index +// CHECK-SAME: %[[VAL_0:.*]]: index) { +// CHECK: %[[VAL_1:.*]] = memref.alloc() : memref<100xf32> +// CHECK: %[[VAL_2:.*]] = memref.alloc() : memref<100xf32, 2> +// CHECK: %[[VAL_3:.*]] = memref.alloc() : memref<1xi32, 4> +// CHECK: %{{.*}} = arith.constant 0 : index +// CHECK: %[[VAL_4:.*]] = arith.constant 64 : index +// CHECK: affine.for %[[VAL_5:.*]] = 0 to 10 { +// CHECK: %[[VAL_6:.*]] = arith.muli %[[VAL_5]], %[[VAL_0]] : index +// CHECK: affine.dma_start %[[VAL_1]]{{\[}}%[[VAL_5]]], %[[VAL_2]]{{\[}}%[[VAL_0]]], %[[VAL_3]]{{\[}}%[[VAL_6]]], %[[VAL_4]] +// CHECK-SAME: : memref<100xf32>, memref<100xf32, 2>, memref<1xi32, 4> +// CHECK: } + +// ----- + +func.func @dma_wait_non_affine_tag_index(%arg0 : index) { + %0 = memref.alloc() : memref<100xf32> + %1 = memref.alloc() : memref<100xf32, 2> + %2 = memref.alloc() : memref<1xi32, 4> + %c0 = arith.constant 0 : index + %c64 = arith.constant 64 : index + affine.for %i0 = 0 to 10 { + %3 = arith.muli %i0, %arg0 : index + affine.dma_wait %2[%3], %c64 : memref<1xi32, 4> + } + return +} + +// CHECK-LABEL: func @dma_wait_non_affine_tag_index +// CHECK-SAME: %[[VAL_0:.*]]: index) { +// CHECK: %{{.*}} = memref.alloc() : memref<100xf32> +// CHECK: %{{.*}} = memref.alloc() : memref<100xf32, 2> +// CHECK: %[[VAL_1:.*]] = memref.alloc() : memref<1xi32, 4> +// CHECK: %{{.*}} = arith.constant 0 : index +// CHECK: %[[VAL_2:.*]] = arith.constant 64 : index +// CHECK: affine.for %[[VAL_3:.*]] = 0 to 10 { +// CHECK: %[[VAL_4:.*]] = arith.muli %[[VAL_3]], %[[VAL_0]] : index +// CHECK: affine.dma_wait %[[VAL_1]]{{\[}}%[[VAL_4]]], %[[VAL_2]] : memref<1xi32, 4> +// CHECK: }