[Codegen][LLVMGPU] Default to private memory space for scalar dispatches (#18523)
Some scalar dispatches can fail to bufferize properly due to
extract_slice foldings. This should be fixed independently, but to at
least compile in such cases, it is fine to always default to private
memory allocations for the LLVMGPUTileAndFuse pipeline since all shared
memory allocations are required to be explicitly managed before
bufferization.
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
index ec74f35..cd89b59 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
@@ -23,6 +23,7 @@
#include "iree/compiler/Dialect/Util/Transforms/Passes.h"
#include "iree/compiler/Utils/PassUtils.h"
#include "llvm/ADT/STLForwardCompat.h"
+#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
#include "mlir/Conversion/ComplexToStandard/ComplexToStandard.h"
@@ -271,18 +272,28 @@
MemRefType memRefType,
ValueRange dynamicSizes,
unsigned alignment) {
- // Bail out if the memref type does not specify a memory space.
- if (!llvm::isa_and_nonnull<gpu::AddressSpaceAttr>(
- memRefType.getMemorySpace())) {
+ Attribute memorySpace = memRefType.getMemorySpace();
+ // Bail out if the memref type specifies a nonnull memory space that is not
+ // #gpu.address_space.
+ if (memorySpace && !llvm::isa<gpu::AddressSpaceAttr>(memorySpace)) {
return failure();
}
+
+ MemRefType allocType = memRefType;
auto privateSpace = gpu::AddressSpaceAttr::get(
builder.getContext(), gpu::GPUDialect::getPrivateAddressSpace());
- if (memRefType.getMemorySpace() == privateSpace) {
- return builder.create<memref::AllocaOp>(loc, memRefType, dynamicSizes)
+ if (!memorySpace) {
+ allocType =
+ MemRefType::get(memRefType.getShape(), memRefType.getElementType(),
+ AffineMap(), privateSpace);
+ memorySpace = privateSpace;
+ }
+
+ if (memorySpace == privateSpace) {
+ return builder.create<memref::AllocaOp>(loc, allocType, dynamicSizes)
.getResult();
}
- return builder.create<memref::AllocOp>(loc, memRefType, dynamicSizes)
+ return builder.create<memref::AllocOp>(loc, allocType, dynamicSizes)
.getResult();
}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_tile_and_fuse.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_tile_and_fuse.mlir
index ab1d896..61c1c24 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_tile_and_fuse.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/pipeline_tile_and_fuse.mlir
@@ -674,3 +674,47 @@
// CHECK-COUNT-32: amdgpu.wmma {{.*}} : vector<16xi8>, vector<16xi8>, vector<8xi32>
// CHECK: vector.transfer_write {{.*}} : vector<4x1x2x8x1xi8>, memref<16x16x196x8x2xi8, #hal.descriptor_type<storage_buffer>>
// CHECK: vector.transfer_write {{.*}} : vector<2x8x1x4x1xi8>, memref<196x8x2x16x16xi8, #hal.descriptor_type<storage_buffer>>
+
+// -----
+
+#lowering_config = #iree_gpu.lowering_config<{
+ thread = [1 : index, 1 : index],
+ workgroup = [1 : index, 1 : index]
+}>
+
+#translation_info = #iree_codegen.translation_info<LLVMGPUTileAndFuse workgroup_size = [32, 1, 1] subgroup_size = 32>
+
+#pipeline_layout = #hal.pipeline.layout<bindings = [
+ #hal.pipeline.binding<storage_buffer, "ReadOnly|Indirect">,
+ #hal.pipeline.binding<storage_buffer, ReadOnly>,
+ #hal.pipeline.binding<storage_buffer, Indirect>
+]>
+
+hal.executable public @main {
+ hal.executable.variant public @rocm_hsaco_fb target(<"rocm", "rocm-hsaco-fb">) {
+ hal.executable.export public @small_elementwise ordinal(0) layout(#pipeline_layout) {
+ ^bb0(%arg0: !hal.device):
+ %x, %y, %z = flow.dispatch.workgroup_count_from_slice
+ hal.return %x, %y, %z : index, index, index
+ }
+ builtin.module {
+ func.func @small_elementwise() attributes {translation_info = #translation_info} {
+ %c0 = arith.constant 0 : index
+ %0 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<1x3xf32>>
+ %1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<1x3xf32>>
+ %2 = hal.interface.binding.subspan layout(#pipeline_layout) binding(2) alignment(64) offset(%c0) flags(Indirect) : !flow.dispatch.tensor<writeonly:tensor<1x3xf32>>
+ %3 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [1, 3], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<1x3xf32>> -> tensor<1x3xf32>
+ %4 = flow.dispatch.tensor.load %1, offsets = [0, 0], sizes = [1, 3], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<1x3xf32>> -> tensor<1x3xf32>
+ %5 = tensor.empty() : tensor<1x3xf32>
+ %6 = linalg.add {lowering_config = #lowering_config} ins(%3, %4 : tensor<1x3xf32>, tensor<1x3xf32>) outs(%5 : tensor<1x3xf32>) -> tensor<1x3xf32>
+ flow.dispatch.tensor.store %6, %2, offsets = [0, 0], sizes = [1, 3], strides = [1, 1] : tensor<1x3xf32> -> !flow.dispatch.tensor<writeonly:tensor<1x3xf32>>
+ return
+ }
+ }
+ }
+}
+
+// CHECK-LABEL: func @small_elementwise
+// CHECK: %[[B:.+]] = hal.interface.binding.subspan layout({{.*}}) binding(2)
+// CHECK: %[[ADD:.+]] = arith.addf %{{.*}}, %{{.*}} : vector<1xf32>
+// CHECK: vector.transfer_write %[[ADD]], %[[B]]
