diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -2348,7 +2348,7 @@ : std::max_element(results.begin(), results.end()); if (resultIt == results.end()) return {}; - return IntegerAttr::get(IndexType::get(op.getContext()), *resultIt); + return Builder(op).getIndexAttr(*resultIt); } /// Remove duplicated expressions in affine min/max ops. @@ -2482,7 +2482,21 @@ // %0 = affine.min (d0) -> (1000, d0 + 512) (%i0) // +static OpFoldResult foldMinZeroOp(AffineMinOp minOp) { + auto isCstZeroExpr = [](AffineExpr expr) { + if (auto cst = expr.dyn_cast()) + return cst.getValue() == 0; + return false; + }; + + if (llvm::any_of(minOp.getAffineMap().getResults(), isCstZeroExpr)) + return Builder(minOp).getIndexAttr(0); + return {}; +} + OpFoldResult AffineMinOp::fold(ArrayRef operands) { + if (auto result = foldMinZeroOp(*this)) + return result; return foldMinMaxOp(*this, operands); } diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir --- a/mlir/test/Dialect/Affine/canonicalize.mlir +++ b/mlir/test/Dialect/Affine/canonicalize.mlir @@ -924,3 +924,13 @@ } return } + +// ----- + +// CHECK-LABEL: func @fold_affine_min_zero +func @fold_affine_min_zero(%i0: index, %i1: index, %i2: index) -> index { + // CHECK: %[[ZERO:.+]] = constant 0 : index + %0 = affine.min affine_map<(d0, d1)[s0] -> (d0, d1 + s0, 0)> (%i0, %i1)[%i2] + // CHECK: return %[[ZERO]] + return %0: index +} diff --git a/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir b/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir --- a/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir +++ b/mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir @@ -411,23 +411,17 @@ return %for0 : tensor } -// CHECK-DAG: #[[DIM2_UNUSED_LOW_MAP:.+]] = affine_map<()[s0] -> (s0, 0)> -// CHECK-DAG: #[[DIM2_USED_LOW_MAP:.+]] = affine_map<()[s0, s1, s2] -> (s0, s1 - s2)> +// CHECK-DAG: #[[DIM2_USED_LOW_MAP:.+]] = affine_map<()[s0, s1] -> (s0, s1)> // CHECK-DAG: #[[DIM2_UNUSED_HIGH_MAP:.+]] = affine_map<()[s0, s1, s2, s3] -> (s0, s0 - s1 + s2 + s3)> -// CHECK-DAG: #[[DIM2_OFFSET_MAP:.+]] = affine_map<()[s0] -> (-s0)> -// CHECK-DAG: #[[DIM3_SIZE_MAP:.+]] = affine_map<()[s0, s1, s2] -> (s0 - s1 - s2)> -// CHECK-DAG: #[[DIM3_UNUSED_HIGH_MAP:.+]] = affine_map<()[s0, s1] -> (0, -s0 + s1)> -// CHECK-DAG: #[[DIM3_USED_HIGH_MAP:.+]] = affine_map<()[s0, s1] -> (s0, -s1)> -// CHECK-DAG: #[[DIM0_UNUSED_LOW_MAP:.+]] = affine_map<(d0)[s0] -> (4, -d0 + s0)> -// CHECK-DAG: #[[DIM0_USED_LOW_MAP:.+]] = affine_map<(d0)[s0] -> (0, 4, -d0 + s0)> -// CHECK-DAG: #[[DIM0_UNUSED_HIGH_MAP:.+]] = affine_map<(d0, d1)[s0] -> (0, -d0 - d1 + s0)> -// CHECK-DAG: #[[DIM0_USED_HIGH_MAP:.+]] = affine_map<(d0, d1)[s0] -> (-d1, 4, -d0 + s0)> -// CHECK-DAG: #[[DIM0_SIZE_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0 - d1 - d2)> +// CHECK-DAG: #[[DIM2_USED_HIGH_MAP:.+]] = affine_map<()[s0, s1, s2] -> (s0, s1 - s2)> +// CHECK-DAG: #[[DIM2_SIZE_MAP:.+]] = affine_map<()[s0, s1, s2] -> (s0 - s1 - s2)> +// CHECK-DAG: #[[DIM0_SIZE_MAP:.+]] = affine_map<(d0)[s0] -> (4, -d0 + s0)> // CHECK-DAG: #[[DIM1_TILE_SIZE_MAP:.+]] = affine_map<(d0)[s0] -> (8, -d0 + s0)> // CHECK-DAG: #[[DIM1_UNUSED_LOW_MAP:.+]] = affine_map<(d0)[s0] -> (s0, d0)> // CHECK-DAG: #[[DIM1_USED_MAP:.+]] = affine_map<(d0, d1)[s0, s1] -> (-d1 + s1, 8, -d0 + s0)> // CHECK-DAG: #[[DIM1_UNUSED_HIGH_MAP:.+]] = affine_map<(d0, d1)[s0, s1, s2] -> (s0, -d0 - d1 + s0 + s1 + s2)> // CHECK-DAG: #[[DIM1_OFFSET_MAP:.+]] = affine_map<(d0, d1) -> (d0 - d1)> +// CHECK-DAG: #[[DIM1_SIZE_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0 - d1 - d2)> // CHECK: func @pad_generic_dynamic // CHECK-SAME: %[[SMALL_INPUT:.+]]: tensor, %[[LARGE_INPUT:.+]]: tensor @@ -442,26 +436,14 @@ // CHECK: %[[LARGE_DIM1:.+]] = memref.dim %[[LARGE_INPUT]], %[[C1]] // CHECK: %[[LARGE_DIM2:.+]] = memref.dim %[[LARGE_INPUT]], %[[C2]] // CHECK: %[[LARGE_DIM3:.+]] = memref.dim %[[LARGE_INPUT]], %[[C3]] -// CHECK: %[[SMALL_DIM0:.+]] = memref.dim %[[SMALL_INPUT]], %[[C0]] // CHECK: %[[SMALL_DIM1:.+]] = memref.dim %[[SMALL_INPUT]], %[[C1]] // CHECK: %[[SMALL_DIM2:.+]] = memref.dim %[[SMALL_INPUT]], %[[C2]] -// CHECK: %[[DIM2_UNUSED_LOW_PAD:.+]] = affine.min #[[DIM2_UNUSED_LOW_MAP]]()[%[[INPUT_LOW_PAD]]] -// CHECK: %[[DIM2_USED_LOW_PAD:.+]] = affine.min #[[DIM2_USED_LOW_MAP]]()[%[[LARGE_DIM2]], %[[INPUT_LOW_PAD]], %[[DIM2_UNUSED_LOW_PAD]]] +// CHECK: %[[DIM2_USED_LOW_PAD:.+]] = affine.min #[[DIM2_USED_LOW_MAP]]()[%[[LARGE_DIM2]], %[[INPUT_LOW_PAD]]] // CHECK: %[[DIM2_UNUSED_HIGH_PAD:.+]] = affine.min #[[DIM2_UNUSED_HIGH_MAP]]()[%[[INPUT_HIGH_PAD]], %[[LARGE_DIM2]], %[[SMALL_DIM2]], %[[INPUT_LOW_PAD]]] -// CHECK: %[[DIM2_USED_HIGH_PAD:.+]] = affine.min #[[DIM2_USED_LOW_MAP]]()[%[[LARGE_DIM2]], %[[INPUT_HIGH_PAD]], %[[DIM2_UNUSED_HIGH_PAD]]] -// CHECK: %[[DIM2_SRC_OFFSET:.+]] = affine.apply #[[DIM2_OFFSET_MAP]]()[%[[DIM2_UNUSED_LOW_PAD]]] -// CHECK: %[[DIM2_SRC_SIZE:.+]] = affine.apply #[[DIM3_SIZE_MAP]]()[%[[LARGE_DIM2]], %[[DIM2_USED_LOW_PAD]], %[[DIM2_USED_HIGH_PAD]]] -// CHECK: %[[SMALL_DIM3:.+]] = memref.dim %[[SMALL_INPUT]], %[[C3]] -// CHECK: %[[DIM3_USED_LOW_PAD:.+]] = affine.min #[[DIM2_UNUSED_LOW_MAP]]()[%[[LARGE_DIM3]]] -// CHECK: %[[DIM3_UNUSED_HIGH_PAD:.+]] = affine.min #[[DIM3_UNUSED_HIGH_MAP]]()[%[[LARGE_DIM3]], %[[SMALL_DIM3]]] -// CHECK: %[[DIM3_USED_HIGH_PAD:.+]] = affine.min #[[DIM3_USED_HIGH_MAP]]()[%[[LARGE_DIM3]], %[[DIM3_UNUSED_HIGH_PAD]]] -// CHECK: %[[DIM3_SRC_SIZE:.+]] = affine.apply #[[DIM3_SIZE_MAP]]()[%[[LARGE_DIM3]], %[[DIM3_USED_LOW_PAD]], %[[DIM3_USED_HIGH_PAD]]] +// CHECK: %[[DIM2_USED_HIGH_PAD:.+]] = affine.min #[[DIM2_USED_HIGH_MAP]]()[%[[LARGE_DIM2]], %[[INPUT_HIGH_PAD]], %[[DIM2_UNUSED_HIGH_PAD]]] +// CHECK: %[[DIM2_SRC_SIZE:.+]] = affine.apply #[[DIM2_SIZE_MAP]]()[%[[LARGE_DIM2]], %[[DIM2_USED_LOW_PAD]], %[[DIM2_USED_HIGH_PAD]]] // CHECK: scf.for %[[IV0:[a-z0-9]+]] -// CHECK: %[[DIM0_UNUSED_LOW_PAD:.+]] = affine.min #[[DIM0_UNUSED_LOW_MAP]](%[[IV0]])[%[[LARGE_DIM0]]] -// CHECK: %[[DIM0_USED_LOW_PAD:.+]] = affine.min #[[DIM0_USED_LOW_MAP]](%[[IV0]])[%[[LARGE_DIM0]]] -// CHECK: %[[DIM0_UNUSED_HIGH_PAD:.+]] = affine.min #[[DIM0_UNUSED_HIGH_MAP]](%[[IV0]], %[[DIM0_UNUSED_LOW_PAD]])[%[[SMALL_DIM0]]] -// CHECK: %[[DIM0_USED_HIGH_PAD:.+]] = affine.min #[[DIM0_USED_HIGH_MAP]](%[[IV0]], %[[DIM0_UNUSED_HIGH_PAD]])[%[[LARGE_DIM0]]] -// CHECK: %[[DIM0_SRC_SIZE:.+]] = affine.apply #[[DIM0_SIZE_MAP]](%[[DIM0_UNUSED_LOW_PAD]], %[[DIM0_USED_LOW_PAD]], %[[DIM0_USED_HIGH_PAD]]) +// CHECK: %[[DIM0_SRC_SIZE:.+]] = affine.min #[[DIM0_SIZE_MAP]](%[[IV0]])[%[[LARGE_DIM0]]] // CHECK: scf.for %[[IV1:[a-z0-9]+]] // CHECK: %[[DIM1_TILE_SIZE:.+]] = affine.min #[[DIM1_TILE_SIZE_MAP]](%[[IV1]])[%[[LARGE_DIM1]]] // CHECK: %[[DIM1_UNUSED_LOW_PAD:.+]] = affine.min #[[DIM1_UNUSED_LOW_MAP]](%[[IV1]])[%[[INPUT_LOW_PAD]]] @@ -469,13 +451,13 @@ // CHECK: %[[DIM1_UNUSED_HIGH_PAD:.+]] = affine.min #[[DIM1_UNUSED_HIGH_MAP]](%[[IV1]], %[[DIM1_TILE_SIZE]])[%[[INPUT_HIGH_PAD]], %[[SMALL_DIM1]], %[[INPUT_LOW_PAD]]] // CHECK: %[[DIM1_USED_HIGH_PAD:.+]] = affine.min #[[DIM1_USED_MAP]](%[[IV1]], %[[DIM1_UNUSED_HIGH_PAD]])[%[[LARGE_DIM1]], %[[INPUT_HIGH_PAD]]] // CHECK: %[[DIM1_SRC_OFFSET:.+]] = affine.apply #[[DIM1_OFFSET_MAP]](%[[IV1]], %[[DIM1_UNUSED_LOW_PAD]]) -// CHECK: %[[DIM1_SRC_SIZE:.+]] = affine.apply #[[DIM0_SIZE_MAP]](%[[DIM1_TILE_SIZE]], %[[DIM1_USED_LOW_PAD]], %[[DIM1_USED_HIGH_PAD]]) +// CHECK: %[[DIM1_SRC_SIZE:.+]] = affine.apply #[[DIM1_SIZE_MAP]](%[[DIM1_TILE_SIZE]], %[[DIM1_USED_LOW_PAD]], %[[DIM1_USED_HIGH_PAD]]) // CHECK: %[[SRC_SUBTENSOR:.+]] = subtensor %[[SMALL_INPUT]] -// CHECK-SAME: [%[[IV0]], %[[DIM1_SRC_OFFSET]], %[[DIM2_SRC_OFFSET]], 0] -// CHECK-SAME: [%[[DIM0_SRC_SIZE]], %[[DIM1_SRC_SIZE]], %[[DIM2_SRC_SIZE]], %[[DIM3_SRC_SIZE]]] +// CHECK-SAME: [%[[IV0]], %[[DIM1_SRC_OFFSET]], 0, 0] +// CHECK-SAME: [%[[DIM0_SRC_SIZE]], %[[DIM1_SRC_SIZE]], %[[DIM2_SRC_SIZE]], %[[LARGE_DIM3]]] // CHECK: %[[PAD:.+]] = linalg.pad_tensor %[[SRC_SUBTENSOR]] -// CHECK-SAME: low[%[[DIM0_USED_LOW_PAD]], %[[DIM1_USED_LOW_PAD]], %[[DIM2_USED_LOW_PAD]], %[[DIM3_USED_LOW_PAD]]] -// CHECK-SAME: high[%[[DIM0_USED_HIGH_PAD]], %[[DIM1_USED_HIGH_PAD]], %[[DIM2_USED_HIGH_PAD]], %[[DIM3_USED_HIGH_PAD]]] +// CHECK-SAME: low[%[[C0]], %[[DIM1_USED_LOW_PAD]], %[[DIM2_USED_LOW_PAD]], %[[C0]]] +// CHECK-SAME: high[%[[C0]], %[[DIM1_USED_HIGH_PAD]], %[[DIM2_USED_HIGH_PAD]], %[[C0]]] // CHECK: linalg.yield %[[ZERO]] // CHECK: linalg.generic // CHECK-SAME: ins(%[[PAD]]