[Codegen][Nearly NFC] Move PropagateDispatchSizeBounds to Common/ (#19650)
As part of refactoring the single-iteration loop remover to use the
ValueBoundsOpInterface, I'll want to annotate workgroup IDs and counts
with theri `upper_bound`s before removing those loops, instead of having
the pass look that information up from context using a custom function.
Since other code like the CPU backend also uses the single-iteration
loop remover and workgroup counts (but not, I think, workitems) I've
generalized the annotation pass to not require a GPU target (and thus to
fall back to not adding a bound if one can't be inferred) and I've moved
it into Codegen/Common (along with renaming it) so it can be called from
non-GPU flows.
diff --git a/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel
index 8582cf9..940477d 100644
--- a/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel
@@ -138,6 +138,7 @@
"PassUtils.cpp",
"Passes.cpp",
"PolynomialApproximationPass.cpp",
+ "PropagateDispatchSizeBounds.cpp",
"PropagateReshapesByExpansion.cpp",
"ReconcileTranslationInfo.cpp",
"RematerializeParallelOps.cpp",
diff --git a/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt
index 1dd9f91..d6e2528 100644
--- a/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt
@@ -130,6 +130,7 @@
"PassUtils.cpp"
"Passes.cpp"
"PolynomialApproximationPass.cpp"
+ "PropagateDispatchSizeBounds.cpp"
"PropagateReshapesByExpansion.cpp"
"ReconcileTranslationInfo.cpp"
"RematerializeParallelOps.cpp"
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel
index f4f8c41..c35b511 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel
@@ -72,7 +72,6 @@
"GPUPatterns.cpp",
"GPUPipelining.cpp",
"GPUPromoteMatmulOperands.cpp",
- "GPUPropagateDispatchSizeBounds.cpp",
"GPUReduceBankConflicts.cpp",
"GPUReuseSharedMemoryAllocs.cpp",
"GPUTensorAlloc.cpp",
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt
index 4c1422e..51576eb 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt
@@ -70,7 +70,6 @@
"GPUPatterns.cpp"
"GPUPipelining.cpp"
"GPUPromoteMatmulOperands.cpp"
- "GPUPropagateDispatchSizeBounds.cpp"
"GPUReduceBankConflicts.cpp"
"GPUReuseSharedMemoryAllocs.cpp"
"GPUTensorAlloc.cpp"
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUPropagateDispatchSizeBounds.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUPropagateDispatchSizeBounds.cpp
deleted file mode 100644
index 43aa70b..0000000
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUPropagateDispatchSizeBounds.cpp
+++ /dev/null
@@ -1,103 +0,0 @@
-// Copyright 2024 The IREE Authors
-//
-// Licensed under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-
-#include "iree/compiler/Codegen/Common/GPU/Passes.h"
-#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
-#include "iree/compiler/Codegen/Utils/GPUUtils.h"
-#include "iree/compiler/Codegen/Utils/Utils.h"
-#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
-#include "mlir/Dialect/GPU/IR/GPUDialect.h"
-#include "mlir/Interfaces/FunctionInterfaces.h"
-#include "mlir/Transforms/Passes.h"
-
-namespace mlir::iree_compiler {
-
-#define GEN_PASS_DEF_GPUPROPAGATEDISPATCHSIZEBOUNDSPASS
-#include "iree/compiler/Codegen/Common/GPU/Passes.h.inc"
-
-namespace {
-
-static void applyBounds(FunctionOpInterface funcOp,
- ArrayRef<int32_t> workgroupSizes,
- ArrayRef<int32_t> workgroupCounts) {
- Builder b(funcOp->getContext());
- funcOp->walk([&](Operation *op) {
- TypeSwitch<Operation *>(op)
- .Case([&](gpu::ThreadIdOp tidOp) {
- tidOp.setUpperBoundAttr(b.getIndexAttr(
- workgroupSizes[static_cast<uint32_t>(tidOp.getDimension())]));
- })
- .Case([&](IREE::HAL::InterfaceWorkgroupSizeOp wgSizeOp) {
- wgSizeOp.setUpperBoundAttr(b.getIndexAttr(
- workgroupSizes[wgSizeOp.getDimension().getZExtValue()]));
- })
- .Case([&](IREE::HAL::InterfaceWorkgroupIDOp wgIdOp) {
- wgIdOp.setUpperBoundAttr(b.getIndexAttr(
- workgroupCounts[wgIdOp.getDimension().getZExtValue()]));
- })
- .Case([&](IREE::HAL::InterfaceWorkgroupCountOp wgCountOp) {
- wgCountOp.setUpperBoundAttr(b.getIndexAttr(
- workgroupCounts[wgCountOp.getDimension().getZExtValue()]));
- })
- .Default([](Operation *) {});
- });
-}
-
-struct GPUPropagateDispatchSizeBoundsPass final
- : impl::GPUPropagateDispatchSizeBoundsPassBase<
- GPUPropagateDispatchSizeBoundsPass> {
- using Base::Base;
-
- void runOnOperation() override {
- FunctionOpInterface funcOp = getOperation();
- IREE::GPU::TargetAttr target = getGPUTargetAttr(funcOp);
- if (!target) {
- funcOp.emitWarning("no known target attribute late in GPU codegen");
- return;
- }
- SmallVector<int32_t, 3> workgroupSizes(
- target.getWgp().getMaxWorkgroupSizes().asArrayRef());
- SmallVector<int32_t, 3> workgroupCounts(
- target.getWgp().getMaxWorkgroupCounts().asArrayRef());
-
- std::optional<SmallVector<int64_t>> staticWorkgroupSize =
- getWorkgroupSize(funcOp);
-
- // Late in codegen, we've reconciled the workgroup size onto the export op.
- if (std::optional<IREE::HAL::ExecutableExportOp> exportOp =
- getEntryPoint(funcOp)) {
- if (std::optional<ArrayAttr> exportWorkgroupSize =
- exportOp->getWorkgroupSize()) {
- staticWorkgroupSize =
- llvm::map_to_vector(exportWorkgroupSize->getAsRange<IntegerAttr>(),
- [](IntegerAttr a) { return a.getInt(); });
- }
- }
-
- if (staticWorkgroupSize) {
- // Target info with no workgroup sizes gives a 0-length array, hence no
- // zip_equal.
- for (auto [size, staticSize] :
- llvm::zip(workgroupSizes, *staticWorkgroupSize)) {
- size = staticSize;
- }
- }
- SmallVector<int64_t> staticWorkgroupCounts = getStaticNumWorkgroups(funcOp);
- assert(staticWorkgroupCounts.size() <= 3 &&
- "workgroup counts are 3D at most");
- for (auto [count, staticCount] :
- llvm::zip(workgroupCounts, staticWorkgroupCounts)) {
- if (staticCount != ShapedType::kDynamic) {
- count = staticCount;
- }
- }
-
- applyBounds(funcOp, workgroupSizes, workgroupCounts);
- }
-};
-} // namespace
-
-} // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td b/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td
index 24552cb..3a71759 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td
@@ -180,11 +180,6 @@
];
}
-def GPUPropagateDispatchSizeBoundsPass :
- InterfacePass<"iree-codegen-gpu-propagate-dispatch-size-bounds", "mlir::FunctionOpInterface"> {
- let summary = "Pass to annotate workitem and workgroup IDs with known bounds";
-}
-
def GPUReduceBankConflictsPass :
InterfacePass<"iree-codegen-gpu-reduce-bank-conflicts", "mlir::FunctionOpInterface"> {
let summary = "Pass to try to reduce the number of bank conflicts by padding memref.alloc ops.";
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel
index 2f3b092..a43b450 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel
@@ -38,7 +38,6 @@
"gpu_pad_operands.mlir",
"gpu_pipeline.mlir",
"gpu_promote_matmul_operands.mlir",
- "gpu_propagate_dispatch_size_bounds.mlir",
"gpu_reorder_workgroups_static.mlir",
"gpu_reorder_workgroups.mlir",
"gpu_reuse_shared_memory_allocs.mlir",
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/GPU/test/CMakeLists.txt
index 50be391..8efb734 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/CMakeLists.txt
@@ -34,7 +34,6 @@
"gpu_pad_operands.mlir"
"gpu_pipeline.mlir"
"gpu_promote_matmul_operands.mlir"
- "gpu_propagate_dispatch_size_bounds.mlir"
"gpu_reorder_workgroups.mlir"
"gpu_reorder_workgroups_static.mlir"
"gpu_reuse_shared_memory_allocs.mlir"
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_propagate_dispatch_size_bounds.mlir b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_propagate_dispatch_size_bounds.mlir
deleted file mode 100644
index f26f2c5..0000000
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_propagate_dispatch_size_bounds.mlir
+++ /dev/null
@@ -1,122 +0,0 @@
-// RUN: iree-opt %s --split-input-file \
-// RUN: --pass-pipeline="builtin.module(hal.executable(hal.executable.variant(builtin.module(func.func(iree-codegen-gpu-propagate-dispatch-size-bounds)))))" \
-// RUN: | FileCheck %s
-
-// Note: not the real target definition, missing types
-#executable_target = #hal.executable.target<"rocm", "rocm-hsaco-fb", {iree.gpu.target = #iree_gpu.target<arch = "gfx1100", features = "",
- wgp = <compute = fp32,
- storage = b32,
- subgroup = arithmetic,
- dot = none, mma = [],
- subgroup_size_choices = [32, 64],
- max_workgroup_sizes = [1024, 1024, 1024],
- max_thread_count_per_workgroup = 1024,
- max_workgroup_memory_bytes = 65536,
- max_workgroup_counts = [2147483647, 2147483647, 2147483647]>>}>
-#pipeline_layout = #hal.pipeline.layout<bindings = [#hal.pipeline.binding<storage_buffer>]>
-
-hal.executable private @static {
- hal.executable.variant public @rocm_hsaco_fb target(#executable_target) {
- hal.executable.export public @static ordinal(0) layout(#pipeline_layout) attributes {workgroup_size = [64 : index, 2 : index, 1 : index]} {
- ^bb0(%arg0: !hal.device):
- %c32 = arith.constant 32 : index
- %c8 = arith.constant 8 : index
- %c1 = arith.constant 1 : index
- hal.return %c32, %c8, %c1 : index, index, index
- }
- builtin.module {
-// CHECK-LABEL: func.func @static
- func.func @static() {
-// CHECK: gpu.thread_id x upper_bound 64
-// CHECK: gpu.thread_id y upper_bound 2
-// CHECK: gpu.thread_id z upper_bound 1
- %thread_id_x = gpu.thread_id x
- %thread_id_y = gpu.thread_id y
- %thread_id_z = gpu.thread_id z
-
-// CHECK: hal.interface.workgroup.size[0] upper_bound 64
-// CHECK: hal.interface.workgroup.size[1] upper_bound 2
-// CHECK: hal.interface.workgroup.size[2] upper_bound 1
- %workgroup_size_x = hal.interface.workgroup.size[0] : index
- %workgroup_size_y = hal.interface.workgroup.size[1] : index
- %workgroup_size_z = hal.interface.workgroup.size[2] : index
-
-// CHECK: hal.interface.workgroup.id[0] upper_bound 32
-// CHECK: hal.interface.workgroup.id[1] upper_bound 8
-// CHECK: hal.interface.workgroup.id[2] upper_bound 1
- %workgroup_id_x = hal.interface.workgroup.id[0] : index
- %workgroup_id_y = hal.interface.workgroup.id[1] : index
- %workgroup_id_z = hal.interface.workgroup.id[2] : index
-
-// CHECK: hal.interface.workgroup.count[0] upper_bound 32
-// CHECK: hal.interface.workgroup.count[1] upper_bound 8
-// CHECK: hal.interface.workgroup.count[2] upper_bound 1
- %workgroup_conut_x = hal.interface.workgroup.count[0] : index
- %workgroup_count_y = hal.interface.workgroup.count[1] : index
- %workgroup_count_z = hal.interface.workgroup.count[2] : index
-
- return
- }
- }
- }
-}
-
-// -----
-
-#executable_target = #hal.executable.target<"rocm", "rocm-hsaco-fb",
- {iree.gpu.target = #iree_gpu.target<arch = "gfx1100", features = "",
- wgp = <compute = fp32,
- storage = b32,
- subgroup = arithmetic,
- dot = none, mma = [],
- subgroup_size_choices = [32, 64],
- max_workgroup_sizes = [1024, 1024, 1024],
- max_thread_count_per_workgroup = 1024,
- max_workgroup_memory_bytes = 65536,
- max_workgroup_counts = [2147483647, 2147483647, 2147483647]>>}>
-#pipeline_layout = #hal.pipeline.layout<bindings = [#hal.pipeline.binding<storage_buffer>]>
-
-hal.executable private @dynamic {
- hal.executable.variant public @rocm_hsaco_fb target(#executable_target) {
- hal.executable.export public @dynamic ordinal(0) layout(#pipeline_layout) {
- ^bb0(%arg0: !hal.device, %arg1: index, %arg2: index):
- %count_x = affine.apply affine_map<()[s0] -> (s0 ceildiv 32)>()[%arg1]
- %count_y = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%arg2]
- %count_z = arith.constant 1 : index
- hal.return %count_x, %count_y, %count_z : index, index, index
- }
- builtin.module {
- func.func @dynamic() {
-// CHECK: gpu.thread_id x upper_bound 1024
-// CHECK: gpu.thread_id y upper_bound 1024
-// CHECK: gpu.thread_id z upper_bound 1024
- %thread_id_x = gpu.thread_id x
- %thread_id_y = gpu.thread_id y
- %thread_id_z = gpu.thread_id z
-
-// CHECK: hal.interface.workgroup.size[0] upper_bound 1024
-// CHECK: hal.interface.workgroup.size[1] upper_bound 1024
-// CHECK: hal.interface.workgroup.size[2] upper_bound 1024
- %workgroup_size_x = hal.interface.workgroup.size[0] : index
- %workgroup_size_y = hal.interface.workgroup.size[1] : index
- %workgroup_size_z = hal.interface.workgroup.size[2] : index
-
-// CHECK: hal.interface.workgroup.id[0] upper_bound 2147483647
-// CHECK: hal.interface.workgroup.id[1] upper_bound 2147483647
-// CHECK: hal.interface.workgroup.id[2] upper_bound 1
- %workgroup_id_x = hal.interface.workgroup.id[0] : index
- %workgroup_id_y = hal.interface.workgroup.id[1] : index
- %workgroup_id_z = hal.interface.workgroup.id[2] : index
-
-// CHECK: hal.interface.workgroup.count[0] upper_bound 2147483647
-// CHECK: hal.interface.workgroup.count[1] upper_bound 2147483647
-// CHECK: hal.interface.workgroup.count[2] upper_bound 1
- %workgroup_conut_x = hal.interface.workgroup.count[0] : index
- %workgroup_count_y = hal.interface.workgroup.count[1] : index
- %workgroup_count_z = hal.interface.workgroup.count[2] : index
-
- return
- }
- }
- }
-}
diff --git a/compiler/src/iree/compiler/Codegen/Common/Passes.td b/compiler/src/iree/compiler/Codegen/Common/Passes.td
index 1854279..7188de2 100644
--- a/compiler/src/iree/compiler/Codegen/Common/Passes.td
+++ b/compiler/src/iree/compiler/Codegen/Common/Passes.td
@@ -531,6 +531,11 @@
let summary = "Convert math operations to their polynomial approximation";
}
+def PropagateDispatchSizeBoundsPass :
+ InterfacePass<"iree-codegen-propagate-dispatch-size-bounds", "mlir::FunctionOpInterface"> {
+ let summary = "Pass to annotate workitem and workgroup IDs with known bounds";
+}
+
def PropagateReshapesByExpansionPass :
Pass<"iree-codegen-propagate-reshapes-by-expansion", ""> {
let summary = "Propagates reshaping operations by expansion.";
diff --git a/compiler/src/iree/compiler/Codegen/Common/PropagateDispatchSizeBounds.cpp b/compiler/src/iree/compiler/Codegen/Common/PropagateDispatchSizeBounds.cpp
new file mode 100644
index 0000000..1ca6ecd
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Common/PropagateDispatchSizeBounds.cpp
@@ -0,0 +1,127 @@
+// Copyright 2024 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Codegen/Common/Passes.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
+#include "iree/compiler/Codegen/Utils/GPUUtils.h"
+#include "iree/compiler/Codegen/Utils/Utils.h"
+#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
+#include "mlir/Transforms/Passes.h"
+
+namespace mlir::iree_compiler {
+
+#define GEN_PASS_DEF_PROPAGATEDISPATCHSIZEBOUNDSPASS
+#include "iree/compiler/Codegen/Common/Passes.h.inc"
+
+namespace {
+
+static void applyBounds(FunctionOpInterface funcOp,
+ ArrayRef<std::optional<int64_t>> workgroupSizes,
+ ArrayRef<std::optional<int64_t>> workgroupCounts) {
+ Builder b(funcOp->getContext());
+ funcOp->walk([&](Operation *op) {
+ TypeSwitch<Operation *>(op)
+ .Case([&](gpu::ThreadIdOp tidOp) {
+ std::optional<int64_t> bound =
+ workgroupSizes[static_cast<uint32_t>(tidOp.getDimension())];
+ if (bound) {
+ tidOp.setUpperBoundAttr(b.getIndexAttr(*bound));
+ }
+ })
+ .Case([&](gpu::BlockDimOp blockDimOp) {
+ std::optional<int64_t> bound =
+ workgroupSizes[static_cast<int32_t>(blockDimOp.getDimension())];
+ if (bound) {
+ blockDimOp.setUpperBoundAttr(b.getIndexAttr(*bound));
+ }
+ })
+ .Case([&](IREE::HAL::InterfaceWorkgroupSizeOp wgSizeOp) {
+ std::optional<int64_t> bound =
+ workgroupSizes[wgSizeOp.getDimension().getZExtValue()];
+ if (bound) {
+ wgSizeOp.setUpperBoundAttr(b.getIndexAttr(*bound));
+ }
+ })
+ .Case([&](IREE::HAL::InterfaceWorkgroupIDOp wgIdOp) {
+ std::optional<int64_t> bound =
+ workgroupCounts[wgIdOp.getDimension().getZExtValue()];
+ if (bound) {
+ wgIdOp.setUpperBoundAttr(b.getIndexAttr(*bound));
+ }
+ })
+ .Case([&](IREE::HAL::InterfaceWorkgroupCountOp wgCountOp) {
+ std::optional<int64_t> bound =
+ workgroupCounts[wgCountOp.getDimension().getZExtValue()];
+ if (bound) {
+ wgCountOp.setUpperBoundAttr(b.getIndexAttr(*bound));
+ }
+ })
+ .Default([](Operation *) {});
+ });
+}
+
+struct PropagateDispatchSizeBoundsPass final
+ : impl::PropagateDispatchSizeBoundsPassBase<
+ PropagateDispatchSizeBoundsPass> {
+ using Base::Base;
+
+ void runOnOperation() override {
+ FunctionOpInterface funcOp = getOperation();
+ SmallVector<std::optional<int64_t>, 3> workgroupSizes(3, std::nullopt);
+ SmallVector<std::optional<int64_t>, 3> workgroupCounts(3, std::nullopt);
+
+ IREE::GPU::TargetAttr target = getGPUTargetAttr(funcOp);
+ if (target) {
+ ArrayRef<int32_t> targetWorkgroupSizes =
+ target.getWgp().getMaxWorkgroupSizes().asArrayRef();
+ ArrayRef<int32_t> targetWorkgroupCounts =
+ target.getWgp().getMaxWorkgroupCounts().asArrayRef();
+ llvm::transform(targetWorkgroupSizes, workgroupSizes.begin(),
+ [](int32_t x) { return std::optional<int64_t>{x}; });
+ llvm::transform(targetWorkgroupCounts, workgroupCounts.begin(),
+ [](int32_t x) { return std::optional<int64_t>{x}; });
+ }
+
+ std::optional<SmallVector<int64_t>> staticWorkgroupSize =
+ getWorkgroupSize(funcOp);
+
+ // Late in codegen, we've reconciled the workgroup size onto the export op.
+ if (std::optional<IREE::HAL::ExecutableExportOp> exportOp =
+ getEntryPoint(funcOp)) {
+ if (std::optional<ArrayAttr> exportWorkgroupSize =
+ exportOp->getWorkgroupSize()) {
+ staticWorkgroupSize =
+ llvm::map_to_vector(exportWorkgroupSize->getAsRange<IntegerAttr>(),
+ [](IntegerAttr a) { return a.getInt(); });
+ }
+ }
+
+ if (staticWorkgroupSize) {
+ // Target info with no workgroup sizes gives a 0-length array, hence no
+ // zip_equal.
+ for (auto [size, staticSize] :
+ llvm::zip(workgroupSizes, *staticWorkgroupSize)) {
+ size = staticSize;
+ }
+ }
+ SmallVector<int64_t> staticWorkgroupCounts = getStaticNumWorkgroups(funcOp);
+ assert(staticWorkgroupCounts.size() <= 3 &&
+ "workgroup counts are 3D at most");
+ for (auto [count, staticCount] :
+ llvm::zip(workgroupCounts, staticWorkgroupCounts)) {
+ if (staticCount != ShapedType::kDynamic) {
+ count = staticCount;
+ }
+ }
+
+ applyBounds(funcOp, workgroupSizes, workgroupCounts);
+ }
+};
+} // namespace
+
+} // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel
index edbb5d8..43a4079 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel
@@ -70,6 +70,7 @@
"optimize_tensor_insert_extract_slices.mlir",
"pad_dynamic_alloc.mlir",
"polynomial_approximation.mlir",
+ "propagate_dispatch_size_bounds.mlir",
"propagate_reshapes_by_expansion.mlir",
"reconcile_translation_info.mlir",
"reductions.mlir",
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt
index e240940..cfeef07 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt
@@ -66,6 +66,7 @@
"optimize_tensor_insert_extract_slices.mlir"
"pad_dynamic_alloc.mlir"
"polynomial_approximation.mlir"
+ "propagate_dispatch_size_bounds.mlir"
"propagate_reshapes_by_expansion.mlir"
"reconcile_translation_info.mlir"
"reductions.mlir"
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/propagate_dispatch_size_bounds.mlir b/compiler/src/iree/compiler/Codegen/Common/test/propagate_dispatch_size_bounds.mlir
new file mode 100644
index 0000000..eb85adc
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Common/test/propagate_dispatch_size_bounds.mlir
@@ -0,0 +1,204 @@
+// RUN: iree-opt %s --split-input-file \
+// RUN: --pass-pipeline="builtin.module(hal.executable(hal.executable.variant(builtin.module(func.func(iree-codegen-propagate-dispatch-size-bounds)))))" \
+// RUN: | FileCheck %s
+
+// Note: not the real target definition, missing types
+#executable_target = #hal.executable.target<"rocm", "rocm-hsaco-fb", {iree.gpu.target = #iree_gpu.target<arch = "gfx1100", features = "",
+ wgp = <compute = fp32,
+ storage = b32,
+ subgroup = arithmetic,
+ dot = none, mma = [],
+ subgroup_size_choices = [32, 64],
+ max_workgroup_sizes = [1024, 1024, 1024],
+ max_thread_count_per_workgroup = 1024,
+ max_workgroup_memory_bytes = 65536,
+ max_workgroup_counts = [2147483647, 2147483647, 2147483647]>>}>
+#pipeline_layout = #hal.pipeline.layout<bindings = [#hal.pipeline.binding<storage_buffer>]>
+
+hal.executable private @static {
+ hal.executable.variant public @rocm_hsaco_fb target(#executable_target) {
+ hal.executable.export public @static ordinal(0) layout(#pipeline_layout) attributes {workgroup_size = [64 : index, 2 : index, 1 : index]} {
+ ^bb0(%arg0: !hal.device):
+ %c32 = arith.constant 32 : index
+ %c8 = arith.constant 8 : index
+ %c1 = arith.constant 1 : index
+ hal.return %c32, %c8, %c1 : index, index, index
+ }
+ builtin.module {
+// CHECK-LABEL: func.func @static()
+ func.func @static() {
+// CHECK: gpu.thread_id x upper_bound 64
+// CHECK: gpu.thread_id y upper_bound 2
+// CHECK: gpu.thread_id z upper_bound 1
+ %thread_id_x = gpu.thread_id x
+ %thread_id_y = gpu.thread_id y
+ %thread_id_z = gpu.thread_id z
+
+// CHECK: gpu.block_dim x upper_bound 64
+// CHECK: gpu.block_dim y upper_bound 2
+// CHECK: gpu.block_dim z upper_bound 1
+ %block_dim_x = gpu.block_dim x
+ %block_dim_y = gpu.block_dim y
+ %block_dim_z = gpu.block_dim z
+
+// CHECK: hal.interface.workgroup.size[0] upper_bound 64
+// CHECK: hal.interface.workgroup.size[1] upper_bound 2
+// CHECK: hal.interface.workgroup.size[2] upper_bound 1
+ %workgroup_size_x = hal.interface.workgroup.size[0] : index
+ %workgroup_size_y = hal.interface.workgroup.size[1] : index
+ %workgroup_size_z = hal.interface.workgroup.size[2] : index
+
+// CHECK: hal.interface.workgroup.id[0] upper_bound 32
+// CHECK: hal.interface.workgroup.id[1] upper_bound 8
+// CHECK: hal.interface.workgroup.id[2] upper_bound 1
+ %workgroup_id_x = hal.interface.workgroup.id[0] : index
+ %workgroup_id_y = hal.interface.workgroup.id[1] : index
+ %workgroup_id_z = hal.interface.workgroup.id[2] : index
+
+// CHECK: hal.interface.workgroup.count[0] upper_bound 32
+// CHECK: hal.interface.workgroup.count[1] upper_bound 8
+// CHECK: hal.interface.workgroup.count[2] upper_bound 1
+ %workgroup_conut_x = hal.interface.workgroup.count[0] : index
+ %workgroup_count_y = hal.interface.workgroup.count[1] : index
+ %workgroup_count_z = hal.interface.workgroup.count[2] : index
+
+ return
+ }
+ }
+ }
+}
+
+// -----
+
+#executable_target = #hal.executable.target<"rocm", "rocm-hsaco-fb",
+ {iree.gpu.target = #iree_gpu.target<arch = "gfx1100", features = "",
+ wgp = <compute = fp32,
+ storage = b32,
+ subgroup = arithmetic,
+ dot = none, mma = [],
+ subgroup_size_choices = [32, 64],
+ max_workgroup_sizes = [1024, 1024, 1024],
+ max_thread_count_per_workgroup = 1024,
+ max_workgroup_memory_bytes = 65536,
+ max_workgroup_counts = [2147483647, 2147483647, 2147483647]>>}>
+#pipeline_layout = #hal.pipeline.layout<bindings = [#hal.pipeline.binding<storage_buffer>]>
+
+hal.executable private @dynamic {
+ hal.executable.variant public @rocm_hsaco_fb target(#executable_target) {
+ hal.executable.export public @dynamic ordinal(0) layout(#pipeline_layout) {
+ ^bb0(%arg0: !hal.device, %arg1: index, %arg2: index):
+ %count_x = affine.apply affine_map<()[s0] -> (s0 ceildiv 32)>()[%arg1]
+ %count_y = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%arg2]
+ %count_z = arith.constant 1 : index
+ hal.return %count_x, %count_y, %count_z : index, index, index
+ }
+ builtin.module {
+// CHECK-LABEL: func.func @dynamic()
+ func.func @dynamic() {
+// CHECK: gpu.thread_id x upper_bound 1024
+// CHECK: gpu.thread_id y upper_bound 1024
+// CHECK: gpu.thread_id z upper_bound 1024
+ %thread_id_x = gpu.thread_id x
+ %thread_id_y = gpu.thread_id y
+ %thread_id_z = gpu.thread_id z
+
+// CHECK: hal.interface.workgroup.size[0] upper_bound 1024
+// CHECK: hal.interface.workgroup.size[1] upper_bound 1024
+// CHECK: hal.interface.workgroup.size[2] upper_bound 1024
+ %workgroup_size_x = hal.interface.workgroup.size[0] : index
+ %workgroup_size_y = hal.interface.workgroup.size[1] : index
+ %workgroup_size_z = hal.interface.workgroup.size[2] : index
+
+// CHECK: hal.interface.workgroup.id[0] upper_bound 2147483647
+// CHECK: hal.interface.workgroup.id[1] upper_bound 2147483647
+// CHECK: hal.interface.workgroup.id[2] upper_bound 1
+ %workgroup_id_x = hal.interface.workgroup.id[0] : index
+ %workgroup_id_y = hal.interface.workgroup.id[1] : index
+ %workgroup_id_z = hal.interface.workgroup.id[2] : index
+
+// CHECK: hal.interface.workgroup.count[0] upper_bound 2147483647
+// CHECK: hal.interface.workgroup.count[1] upper_bound 2147483647
+// CHECK: hal.interface.workgroup.count[2] upper_bound 1
+ %workgroup_conut_x = hal.interface.workgroup.count[0] : index
+ %workgroup_count_y = hal.interface.workgroup.count[1] : index
+ %workgroup_count_z = hal.interface.workgroup.count[2] : index
+
+ return
+ }
+ }
+ }
+}
+
+// -----
+
+#executable_target = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu_features = "+avx512f", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "x86_64-unknown-linux-gnu"}>
+#pipeline_layout = #hal.pipeline.layout<bindings = [#hal.pipeline.binding<storage_buffer>]>
+
+hal.executable private @static_cpu {
+ hal.executable.variant public @embedded_elf_x86_64 target(#executable_target) {
+ hal.executable.export public @static_cpu ordinal(0) layout(#pipeline_layout) attributes {workgroup_size = [64 : index, 2 : index, 1 : index]} {
+ ^bb0(%arg0: !hal.device):
+ %c32 = arith.constant 32 : index
+ %c8 = arith.constant 8 : index
+ %c1 = arith.constant 1 : index
+ hal.return %c32, %c8, %c1 : index, index, index
+ }
+ builtin.module {
+// CHECK-LABEL: func.func @static_cpu()
+ func.func @static_cpu() {
+// CHECK: hal.interface.workgroup.id[0] upper_bound 32
+// CHECK: hal.interface.workgroup.id[1] upper_bound 8
+// CHECK: hal.interface.workgroup.id[2] upper_bound 1
+ %workgroup_id_x = hal.interface.workgroup.id[0] : index
+ %workgroup_id_y = hal.interface.workgroup.id[1] : index
+ %workgroup_id_z = hal.interface.workgroup.id[2] : index
+
+// CHECK: hal.interface.workgroup.count[0] upper_bound 32
+// CHECK: hal.interface.workgroup.count[1] upper_bound 8
+// CHECK: hal.interface.workgroup.count[2] upper_bound 1
+ %workgroup_conut_x = hal.interface.workgroup.count[0] : index
+ %workgroup_count_y = hal.interface.workgroup.count[1] : index
+ %workgroup_count_z = hal.interface.workgroup.count[2] : index
+
+ return
+ }
+ }
+ }
+}
+
+// -----
+
+#executable_target = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu_features = "+avx512f", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "x86_64-unknown-linux-gnu"}>
+#pipeline_layout = #hal.pipeline.layout<bindings = [#hal.pipeline.binding<storage_buffer>]>
+
+hal.executable private @dynamic_cpu {
+ hal.executable.variant public @embedded_elf_x86_64 target(#executable_target) {
+ hal.executable.export public @dynamic_cpu ordinal(0) layout(#pipeline_layout) {
+ ^bb0(%arg0: !hal.device, %arg1: index, %arg2: index):
+ %count_x = affine.apply affine_map<()[s0] -> (s0 ceildiv 32)>()[%arg1]
+ %count_y = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%arg2]
+ %count_z = arith.constant 1 : index
+ hal.return %count_x, %count_y, %count_z : index, index, index
+ }
+ builtin.module {
+// CHECK-LABEL: @dynamic_cpu()
+ func.func @dynamic_cpu() {
+// CHECK: hal.interface.workgroup.id[0] : index
+// CHECK: hal.interface.workgroup.id[1] : index
+// CHECK: hal.interface.workgroup.id[2] upper_bound 1 : index
+ %workgroup_id_x = hal.interface.workgroup.id[0] : index
+ %workgroup_id_y = hal.interface.workgroup.id[1] : index
+ %workgroup_id_z = hal.interface.workgroup.id[2] : index
+
+// CHECK: hal.interface.workgroup.count[0] : index
+// CHECK: hal.interface.workgroup.count[1] : index
+// CHECK: hal.interface.workgroup.count[2] upper_bound 1 : index
+ %workgroup_conut_x = hal.interface.workgroup.count[0] : index
+ %workgroup_count_y = hal.interface.workgroup.count[1] : index
+ %workgroup_count_z = hal.interface.workgroup.count[2] : index
+
+ return
+ }
+ }
+ }
+}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
index d0a269e..c3dda57 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
@@ -1053,7 +1053,7 @@
.addPass(createLLVMGPUVectorLoweringPass)
.addPass(createExpandGPUOpsPass)
// Expose workitem and workgroup counts to range inference later.
- .addPass(createGPUPropagateDispatchSizeBoundsPass);
+ .addPass(createPropagateDispatchSizeBoundsPass);
// This pass needs to run before SCF -> CF.
addLowerAndOptimizeAddressComputationPasses(funcPassManager);
diff --git a/compiler/src/iree/compiler/Codegen/SPIRV/Passes.cpp b/compiler/src/iree/compiler/Codegen/SPIRV/Passes.cpp
index 511dbe7..2a13558 100644
--- a/compiler/src/iree/compiler/Codegen/SPIRV/Passes.cpp
+++ b/compiler/src/iree/compiler/Codegen/SPIRV/Passes.cpp
@@ -227,7 +227,7 @@
/// Adds passes to perform the final SPIR-V conversion.
static void addSPIRVLoweringPasses(OpPassManager &modulePassManager) {
FunctionLikeNest(modulePassManager)
- .addPass(createGPUPropagateDispatchSizeBoundsPass)
+ .addPass(createPropagateDispatchSizeBoundsPass)
.addPass(createCanonicalizerPass)
.addPass(createCSEPass)
.addPass(createLowerAffinePass)
