Skip to content

Commit 391a757

Browse files
authored
[mlir][gpu] Add lowering dynamic_shared_memory op for rocdl (#74473)
This PR adds lowering of `gpu.dynamic_shared_memory` to rocdl target.
1 parent 9c2e544 commit 391a757

File tree

2 files changed

+19
-0
lines changed

2 files changed

+19
-0
lines changed

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -363,6 +363,8 @@ void mlir::populateGpuToROCDLConversionPatterns(
363363
// Use address space = 4 to match the OpenCL definition of printf()
364364
patterns.add<GPUPrintfOpToLLVMCallLowering>(converter, /*addressSpace=*/4);
365365
}
366+
// TODO: Add alignment for workgroup memory
367+
patterns.add<GPUDynamicSharedMemoryOpLowering>(converter);
366368

367369
patterns.add<GPUShuffleOpLowering, GPULaneIdOpToROCDL>(converter);
368370

mlir/test/Conversion/GPUCommon/lower-memory-space-attrs.mlir

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,3 +46,20 @@ gpu.module @kernel {
4646
// CHECK: [[value:%.+]] = llvm.load
4747
// CHECK-SAME: : !llvm.ptr<1> -> f32
4848
// CHECK: llvm.return [[value]]
49+
50+
// -----
51+
52+
gpu.module @kernel {
53+
gpu.func @dynamic_shmem(%arg0: f32) {
54+
%0 = arith.constant 0 : index
55+
%1 = gpu.dynamic_shared_memory : memref<?xi8, #gpu.address_space<workgroup>>
56+
%2 = memref.view %1[%0][] : memref<?xi8, #gpu.address_space<workgroup>> to memref<4xf32, #gpu.address_space<workgroup>>
57+
memref.store %arg0, %2[%0] : memref<4xf32, #gpu.address_space<workgroup>>
58+
gpu.return
59+
}
60+
}
61+
62+
// CHECK-LABEL: llvm.func @dynamic_shmem
63+
// CHECK: llvm.store
64+
// CHECK-SAME: : f32, !llvm.ptr<3>
65+

0 commit comments

Comments
 (0)