[Mlir-commits] [mlir] 391a757 - [mlir][gpu] Add lowering dynamic_shared_memory op for rocdl (#74473)

[llvmlistbot at llvm.org]

Author: Guray Ozen
Date: 2023-12-05T19:56:43+01:00
New Revision: 391a7577e703516dbefd41b3da8f3bbd751c6978

URL: https://github.com/llvm/llvm-project/commit/391a7577e703516dbefd41b3da8f3bbd751c6978
DIFF: https://github.com/llvm/llvm-project/commit/391a7577e703516dbefd41b3da8f3bbd751c6978.diff

LOG: [mlir][gpu] Add lowering dynamic_shared_memory op for rocdl (#74473)

This PR adds lowering of `gpu.dynamic_shared_memory` to rocdl target.

Added: 
    

Modified: 
    mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
    mlir/test/Conversion/GPUCommon/lower-memory-space-attrs.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
index d9f94e30b04c6..599bb13190f12 100644
--- a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
+++ b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
@@ -363,6 +363,8 @@ void mlir::populateGpuToROCDLConversionPatterns(
     // Use address space = 4 to match the OpenCL definition of printf()
     patterns.add<GPUPrintfOpToLLVMCallLowering>(converter, /*addressSpace=*/4);
   }
+  // TODO: Add alignment for workgroup memory
+  patterns.add<GPUDynamicSharedMemoryOpLowering>(converter);
 
   patterns.add<GPUShuffleOpLowering, GPULaneIdOpToROCDL>(converter);
 

diff  --git a/mlir/test/Conversion/GPUCommon/lower-memory-space-attrs.mlir b/mlir/test/Conversion/GPUCommon/lower-memory-space-attrs.mlir
index 14f5302ac2002..f8d82f533cc0b 100644
--- a/mlir/test/Conversion/GPUCommon/lower-memory-space-attrs.mlir
+++ b/mlir/test/Conversion/GPUCommon/lower-memory-space-attrs.mlir
@@ -46,3 +46,20 @@ gpu.module @kernel {
 //       CHECK: [[value:%.+]] = llvm.load
 //  CHECK-SAME:   : !llvm.ptr<1> -> f32
 //       CHECK: llvm.return [[value]]
+
+// -----
+
+gpu.module @kernel {
+  gpu.func @dynamic_shmem(%arg0: f32)  {
+    %0 = arith.constant 0 : index
+    %1 = gpu.dynamic_shared_memory : memref<?xi8, #gpu.address_space<workgroup>>
+    %2 = memref.view %1[%0][] : memref<?xi8, #gpu.address_space<workgroup>> to memref<4xf32, #gpu.address_space<workgroup>>
+    memref.store %arg0, %2[%0] : memref<4xf32, #gpu.address_space<workgroup>>
+    gpu.return
+  }
+}
+
+// CHECK-LABEL:  llvm.func @dynamic_shmem
+//       CHECK:  llvm.store
+//  CHECK-SAME:   : f32, !llvm.ptr<3>
+