[Codegen] Replace LLVMGPUVectorize with LLVMGPUTileAndFuse for supported cases (#18474)
The plan is to deprecate the LLVMGPUVectorize pipeline as it is no
longer maintained. This patch is the first step towards deprecating it.
The pipeline is still used by default for convolutions (if
LLVMGPUVectorDistribute does not pick it up).
1) (this) Replace default use cases of LLVMGPUVectorize with
LLVMGPUTileAndFuse.
2) Replace convolution focused uses of LLVMGPUVectorize.
3) Separate out the lowering config logic for pre-bufferized dispatches
(that targets LLVMGPUBaseDistribute).
4) Delete the vectorize pipeline.
diff --git a/compiler/plugins/target/CUDA/test/smoketest.mlir b/compiler/plugins/target/CUDA/test/smoketest.mlir
index 0c12f06..0ac7a8e 100644
--- a/compiler/plugins/target/CUDA/test/smoketest.mlir
+++ b/compiler/plugins/target/CUDA/test/smoketest.mlir
@@ -62,11 +62,11 @@
}
// PTX: .entry add_dispatch
-// PTX: .maxntid 64, 1, 1
+// PTX: .maxntid 32, 1, 1
// PTX: add.rn.f32
// PTX: .entry mul_dispatch
-// PTX: .maxntid 64, 1, 1
+// PTX: .maxntid 32, 1, 1
// PTX: mul.rn.f32
// CHECK: hal.executable public @smoketest_linked
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_apply_tiling_level.mlir b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_apply_tiling_level.mlir
index 3350078..841a79c 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_apply_tiling_level.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_apply_tiling_level.mlir
@@ -523,3 +523,37 @@
// NOZERO: tensor.pad
// NOZERO: tensor.pad
// NOZERO: linalg.matmul
+
+// -----
+
+func.func @distribute_multi_result_generic(
+ %arg0: tensor<3x4x5xf32>, %arg1: tensor<3x4xf32>, %arg2: tensor<3x4xf32>) -> (tensor<3x4x5xf32>, tensor<3x4x5xf32>)
+ attributes {
+ translation_info = #iree_codegen.translation_info<LLVMGPUTileAndFuse workgroup_size = [32, 1, 1] subgroup_size = 32, {}>
+ } {
+ %empty = tensor.empty() : tensor<3x4x5xf32>
+ %0:2 = linalg.generic {
+ indexing_maps = [
+ affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
+ affine_map<(d0, d1, d2) -> (d0, d1)>,
+ affine_map<(d0, d1, d2) -> (d0, d1)>,
+ affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
+ affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+ ], iterator_types = ["parallel", "parallel", "parallel"]}
+ ins(%arg0, %arg1, %arg2 : tensor<3x4x5xf32>, tensor<3x4xf32>, tensor<3x4xf32>)
+ outs(%empty, %empty : tensor<3x4x5xf32>, tensor<3x4x5xf32>)
+ attrs = {lowering_config = #iree_gpu.derived_thread_config} {
+ ^bb0(%in: f32, %in_1: f32, %in_2: f32, %out: f32, %out_3: f32):
+ %7 = arith.subf %in, %in_1 : f32
+ %8 = math.exp %7 : f32
+ %9 = arith.divf %8, %in_2 : f32
+ linalg.yield %8, %9 : f32, f32
+ } -> (tensor<3x4x5xf32>, tensor<3x4x5xf32>)
+ return %0#0, %0#1 : tensor<3x4x5xf32>, tensor<3x4x5xf32>
+}
+
+// THREAD-LABEL: @distribute_multi_result_generic
+// THREAD: %[[FORALL:.+]]:2 = scf.forall
+// THREAD: linalg.generic
+// THREAD-SAME: outs(%{{.*}}, %{{.*}}: tensor<1x1x?xf32>, tensor<1x1x?xf32>)
+// THREAD: return %[[FORALL]]#0, %[[FORALL]]#1
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/DerivedConfigUtils.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/DerivedConfigUtils.cpp
index 4e534d4..6f7b75f 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/DerivedConfigUtils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/DerivedConfigUtils.cpp
@@ -106,10 +106,6 @@
if (!linalgOp.hasPureTensorSemantics()) {
return {};
}
- // TODO: Support multi-result
- if (linalgOp->getNumResults() != 1) {
- return {};
- }
SmallVector<int64_t> loopRanges = linalgOp.getStaticLoopRanges();
int64_t vectorSize = kPreferredCopyNumBits /
getElementTypeOrSelf(linalgOp->getResultTypes()[0])
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp
index 1a8697c..4ee5024 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp
@@ -622,7 +622,8 @@
};
// First try to see if we can use up all threads without any loss.
- if (distributeToThreads(subgroupSize) != 1) {
+ int64_t newNumThreads = subgroupSize;
+ if (distributeToThreads(newNumThreads) != 1) {
// Otherwise, allow larger and larger loss factor.
// Threads for distribution. Use 32 at least.
@@ -636,22 +637,6 @@
}
}
- // TODO(qedawkins): Currently scf.forall resolution only supports static
- // trip counts, meaning the workgroup tile size must perfectly divide the
- // loop bound (and thread tile size must perfectly divide the workgroup tile)
- // so that the trip count won't be static. Remove this check once proper
- // dynamic trip count resolution support is added.
- for (auto [loopId, threadTile] : llvm::enumerate(threadTileSizes)) {
- if (threadTile == 0) {
- continue;
- }
- int64_t bound = loopBounds[loopId];
- int64_t wkgpTile = workgroupTileSizes[loopId];
- if (bound % wkgpTile != 0 || wkgpTile % threadTile != 0) {
- return failure();
- }
- }
-
// Attach the MMA schedule as an attribute to the entry point export function
// for later access in the pipeline.
MLIRContext *context = linalgOp.getContext();
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
index f269e2b..7fc5b80 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
@@ -57,6 +57,11 @@
"test the tile and fuse pipeline for all supported operations"),
llvm::cl::init(false));
+llvm::cl::opt<bool> clLLVMGPUVectorizePipeline(
+ "iree-codegen-llvmgpu-vectorize-pipeline",
+ llvm::cl::desc("forces use of the legacy LLVMGPU vectorize pipeline"),
+ llvm::cl::init(false));
+
llvm::cl::opt<bool> clGPUEnableVectorDistribution(
"iree-codegen-llvmgpu-use-vector-distribution",
llvm::cl::desc("enable the usage of the vector distribution pipeline"),
@@ -2175,6 +2180,13 @@
})
.Default([&](auto op) {
LDBG("Default Config");
+ if (!clLLVMGPUVectorizePipeline) {
+ if (succeeded(IREE::GPU::setTileAndFuseLoweringConfig(
+ target, entryPointFn, computeOp))) {
+ LDBG("Tile and fuse default config");
+ return success();
+ }
+ }
return setRootDefaultConfig(target, entryPointFn, computeOp);
});
}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
index d13a5a4..53f67e4 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
@@ -386,8 +386,6 @@
funcPassManager.addPass(createPropagateReshapesByExpansionPass());
funcPassManager.addPass(createCanonicalizerPass());
funcPassManager.addPass(createCSEPass());
- funcPassManager.addPass(createConvertToDestinationPassingStylePass(
- /*useWARForCooperativeMatrixCodegen=*/false));
// Step 4. Tile and fuse tileable ops to subgroups/threads.
{
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_tile_and_fuse.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_tile_and_fuse.mlir
index 80d4691..9c75196 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_tile_and_fuse.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_tile_and_fuse.mlir
@@ -206,10 +206,8 @@
}
}
-// Verify that this does not select this pipeline due to issues with resolving
-// dynamic scf.forall loops.
-// CHECK-LABEL: module @elementwise_unaligned
-// CHECK-NOT: LLVMGPUTileAndFuse
+// CHECK-LABEL: func.func @elementwise_unaligned
+// CHECK-SAME: #iree_codegen.translation_info<LLVMGPUTileAndFuse workgroup_size = [64, 1, 1] subgroup_size = 64>
// -----
@@ -225,7 +223,7 @@
// Verify that a lowering config is set on large rank tensors with unaligned
// shapes.
// CHECK-LABEL: func.func @elementwise_large_rank
-// CHECK-SAME: #iree_codegen.translation_info<LLVMGPUVectorize workgroup_size = [128, 1, 1] subgroup_size = 64>
+// CHECK-SAME: #iree_codegen.translation_info<LLVMGPUTileAndFuse workgroup_size = [64, 1, 1] subgroup_size = 64>
// -----
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/elementwise_pipeline.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/elementwise_pipeline.mlir
index 7f50d1d..4fd68e7 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/elementwise_pipeline.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/elementwise_pipeline.mlir
@@ -13,7 +13,8 @@
%1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) alignment(64) offset(%c0) : !flow.dispatch.tensor<writeonly:tensor<320x320x3x3xf32>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [3, 320, 320, 3], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<3x320x320x3xf32>> -> tensor<3x320x320x3xf32>
%3 = tensor.empty() : tensor<320x320x3x3xf32>
- %4 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%2 : tensor<3x320x320x3xf32>) outs(%3 : tensor<320x320x3x3xf32>) {
+ %4 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
+ ins(%2 : tensor<3x320x320x3xf32>) outs(%3 : tensor<320x320x3x3xf32>) {
^bb0(%in: f32, %out: f32):
%5 = arith.addf %in, %cst : f32
linalg.yield %5 : f32
@@ -21,6 +22,6 @@
flow.dispatch.tensor.store %4, %1, offsets = [0, 0, 0, 0], sizes = [320, 320, 3, 3], strides = [1, 1, 1, 1] : tensor<320x320x3x3xf32> -> !flow.dispatch.tensor<writeonly:tensor<320x320x3x3xf32>>
return
}
-// CHECK: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<LLVMGPUVectorize workgroup_size = [3, 3, 7] subgroup_size = 32>
+// CHECK: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<LLVMGPUTileAndFuse workgroup_size = [32, 1, 1] subgroup_size = 32>
// CHECK: func.func @forward_dispatch_0_generic_320x320x3x3()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/gpu_set_num_workgroups.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/gpu_set_num_workgroups.mlir
index 43a7164..c56b0e7 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/gpu_set_num_workgroups.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/gpu_set_num_workgroups.mlir
@@ -28,12 +28,11 @@
return
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[256]{{\]}}>
-// CHECK: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<LLVMGPUVectorize workgroup_size = [64, 1, 1] subgroup_size = 32>
+// CHECK: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<LLVMGPUTileAndFuse workgroup_size = [32, 1, 1] subgroup_size = 32>
// CHECK: func.func @add_dispatch_0
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: linalg.generic
-// CHECK-SAME: lowering_config = #[[CONFIG]]
+// CHECK-SAME: lowering_config = #iree_gpu.lowering_config<{thread = [4], workgroup = [128]}>
// -----
@@ -153,14 +152,15 @@
return
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[1, 128, 4]{{\]}}>
-// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<LLVMGPUVectorize workgroup_size = [64, 1, 1] subgroup_size = 32>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<LLVMGPUTileAndFuse workgroup_size = [32, 1, 1] subgroup_size = 32>
// CHECK: func.func @reduction_aligned2()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: linalg.fill
-// CHECK-SAME: lowering_config = #[[CONFIG]]
// CHECK: linalg.generic
-// CHECK-SAME: lowering_config = #[[CONFIG]]
+// CHECK-SAME: lowering_config = #iree_gpu.lowering_config
+// CHECK-SAME: reduction = [0, 0, 4]
+// CHECK-SAME: thread = [1, 4, 0]
+// CHECK-SAME: workgroup = [1, 128, 0]
// -----
@@ -436,7 +436,7 @@
return
}
-// SM80-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<LLVMGPUVectorize workgroup_size = [64, 1, 1] subgroup_size = 32
+// SM80-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<LLVMGPUTileAndFuse workgroup_size = [32, 1, 1] subgroup_size = 32
// SM80: func.func @contract_reduction()
// SM80-SAME: translation_info = #[[TRANSLATION]]
@@ -563,12 +563,11 @@
return
}
-// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[64]{{\]}}>
-// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<LLVMGPUVectorize workgroup_size = [64, 1, 1] subgroup_size = 32>
+// CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<LLVMGPUTileAndFuse workgroup_size = [32, 1, 1] subgroup_size = 32>
// CHECK: func.func @inner_unit_dim()
// CHECK-SAME: translation_info = #[[TRANSLATION]]
// CHECK: linalg.generic
-// CHECK-SAME: lowering_config = #[[CONFIG]]
+// CHECK-SAME: lowering_config = #iree_gpu.lowering_config<{thread = [1, 1], workgroup = [32, 1]}>
// -----
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/nvvm_pipeline_test.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/nvvm_pipeline_test.mlir
index c591885..9bb0e90 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/nvvm_pipeline_test.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/nvvm_pipeline_test.mlir
@@ -861,22 +861,24 @@
#hal.pipeline.binding<storage_buffer>,
#hal.pipeline.binding<storage_buffer>
]>
-hal.executable private @shared_mem_alloc {
+hal.executable private @generalized_pool {
hal.executable.variant public @cuda_nvptx_fb target(<"cuda", "cuda-nvptx-fb">) {
- hal.executable.export public @shared_mem_alloc ordinal(0) layout(#pipeline_layout) {
+ hal.executable.export public @generalized_pool ordinal(0) layout(#pipeline_layout) {
^bb0(%arg0: !hal.device, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index):
%x, %y, %z = flow.dispatch.workgroup_count_from_dag_root %arg1, %arg2, %arg3, %arg4, %arg5
hal.return %x, %y, %z : index, index, index
}
builtin.module {
- func.func @shared_mem_alloc() {
+ func.func @generalized_pool() {
%c0 = arith.constant 0 : index
- %cst = arith.constant dense<0xFF800000> : tensor<14x14x480xf32>
+ %cst = arith.constant 0xFF800000 : f32
+ %empty = tensor.empty() : tensor<14x14x480xf32>
+ %fill = linalg.fill ins(%cst : f32) outs(%empty : tensor<14x14x480xf32>) -> tensor<14x14x480xf32>
%0 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0) : !flow.dispatch.tensor<readonly:tensor<29x29x480xf32>>
%1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) alignment(64) offset(%c0) : !flow.dispatch.tensor<writeonly:tensor<14x14x480xf32>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0], sizes = [29, 29, 480], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<29x29x480xf32>> -> tensor<29x29x480xf32>
%3 = tensor.empty() : tensor<3x3xf32>
- %4 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d0 * 2 + d3, d1 * 2 + d4, d2)>, affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%2, %3 : tensor<29x29x480xf32>, tensor<3x3xf32>) outs(%cst : tensor<14x14x480xf32>) {
+ %4 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d0 * 2 + d3, d1 * 2 + d4, d2)>, affine_map<(d0, d1, d2, d3, d4) -> (d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"]} ins(%2, %3 : tensor<29x29x480xf32>, tensor<3x3xf32>) outs(%fill : tensor<14x14x480xf32>) {
^bb0(%arg0: f32, %arg1: f32, %arg2: f32):
%5 = arith.maximumf %arg2, %arg0 : f32
linalg.yield %5 : f32
@@ -888,16 +890,11 @@
}
}
-// Check that bufferization is emitting correct code for the temp shared
-// memory alloc.
-// CHECK-LABEL: hal.executable private @shared_mem_alloc
+// CHECK-LABEL: hal.executable private @generalized_pool
// CHECK: hal.executable.variant public @cuda
-// CHECK: nvvm.barrier0
-// CHECK: llvm.store %{{.*}}, %{{.*}} : f32, !llvm.ptr<3>
-// CHECK: nvvm.barrier0
-// CHECK: nvvm.barrier0
-// CHECK: llvm.load %{{.*}} : !llvm.ptr<3> -> f32
-// CHECK: nvvm.barrier0
+// CHECK: llvm.load %{{.*}} : !llvm.ptr<1> -> f32
+// CHECK: llvm.intr.maxnum
+// CHECK: llvm.store %{{.*}}, %{{.*}} : f32, !llvm.ptr<1>
// -----
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_cuda.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_cuda.mlir
index e7aaae0..63d07a3 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_cuda.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/reduction_pipeline_cuda.mlir
@@ -276,9 +276,8 @@
// Small reduction computes the whole reduction on a single thread.
// CHECK-LABEL: func.func @small_reduction
-// CHECK: scf.for %{{.*}} = %c0 to %c13 step %c4
-// CHECK: linalg.generic
-// CHECK: arith.addf
+// CHECK: %[[READ:.+]] = vector.transfer_read {{.*}} #hal.descriptor_type<storage_buffer>>, vector<4x13xf32>
+// CHECK: vector.multi_reduction <add>, %[[READ]], {{.*}} : vector<4x13xf32> to vector<4xf32>
// -----
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/rocdl_pipeline_test.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/rocdl_pipeline_test.mlir
index 7736a40..edddd19 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/rocdl_pipeline_test.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/rocdl_pipeline_test.mlir
@@ -2,7 +2,7 @@
// RUN: iree-opt --split-input-file --iree-gpu-test-target=gfx940 --pass-pipeline="builtin.module(hal.executable(hal.executable.variant(builtin.module(iree-codegen-llvmgpu-configuration-pipeline), iree-codegen-linalg-to-rocdl-pipeline)))" %s | FileCheck %s --check-prefix=CDNA3
// RUN: iree-opt --split-input-file --iree-gpu-test-target=gfx1100 --pass-pipeline="builtin.module(hal.executable(hal.executable.variant(builtin.module(iree-codegen-llvmgpu-configuration-pipeline), iree-codegen-linalg-to-rocdl-pipeline)))" %s | FileCheck %s --check-prefix=RDNA3
-// Verify that a simple element wise op gets lowered succefully all the way to
+// Verify that a simple element wise op gets lowered successfully all the way to
// nvvm/llvm dialect.
#pipeline_layout = #hal.pipeline.layout<bindings = [
@@ -142,7 +142,7 @@
// CDNA3-LABEL: hal.executable public @ext_fp8_dispatch
// CDNA3: hal.executable.variant public @rocm
-// CDNA3-DAG: %[[EXTF8E4M3:.+]] = rocdl.cvt.f32.fp8 %{{.*}} : f32
-// CDNA3-DAG: %[[EXTF8E5M2:.+]] = rocdl.cvt.f32.bf8 %{{.*}} : f32
-// CDNA3: %[[ADD:.+]] = llvm.fadd %[[EXTF8E4M3]], %[[EXTF8E5M2]] : f32
-// CDNA3: llvm.store %[[ADD]], %{{.*}} : f32, !llvm.ptr<1>
+// CDNA3-COUNT-4: rocdl.cvt.f32.fp8 %{{.*}} : f32
+// CDNA3-COUNT-4: rocdl.cvt.f32.bf8 %{{.*}} : f32
+// CDNA3: %[[ADD:.+]] = llvm.fadd %{{.*}}, %{{.*}} : vector<4xf32>
+// CDNA3: llvm.store %[[ADD]], %{{.*}} : vector<4xf32>, !llvm.ptr<1>
diff --git a/compiler/src/iree/compiler/Codegen/Transforms/Transforms.cpp b/compiler/src/iree/compiler/Codegen/Transforms/Transforms.cpp
index 6d0d052..5fe45e9 100644
--- a/compiler/src/iree/compiler/Codegen/Transforms/Transforms.cpp
+++ b/compiler/src/iree/compiler/Codegen/Transforms/Transforms.cpp
@@ -1545,8 +1545,6 @@
// Step 5. Inline the body of the original forall into the new for loop.
OpBuilder::InsertionGuard g2(rewriter);
SmallVector<Value> argReplacements(newForallOp.getInductionVars());
- // If there is no paired slice (for a single trip count loop) then
- // replace the iter arg of the forall op directly.
for (auto [forallIterArg, forIterArg, maybeSlice] :
llvm::zip_equal(newForallOp.getRegionIterArgs(),
newLoop.getRegionIterArgs(), pairedSlices)) {
