[Mlir-commits] [mlir] [mlir][gpu] Introduce `gpu.dynamic_shared_memory` Op (PR #71546)

[Guray Ozen]

================
@@ -554,6 +555,95 @@ static IntegerAttr wrapNumericMemorySpace(MLIRContext *ctx, unsigned space) {
   return IntegerAttr::get(IntegerType::get(ctx, 64), space);
 }
 
+/// Generates a symbol with 0-sized array type for dynamic shared memory usage,
+/// or uses existing symbol.
+LLVM::GlobalOp
+getDynamicSharedMemorySymbol(ConversionPatternRewriter &rewriter,
+                             gpu::DynamicSharedMemoryOp op,
+                             const LLVMTypeConverter *typeConverter,
+                             MemRefType memrefType, unsigned alignmentBit) {
+  std::optional<LLVM::GlobalOp> existingGlobalOp;
+
+  LLVM::LLVMFuncOp funcOp = op->getParentOfType<LLVM::LLVMFuncOp>();
+  assert(funcOp && "cannot find llvm.func op");
+
+  gpu::GPUModuleOp moduleOp = funcOp->getParentOfType<gpu::GPUModuleOp>();
+  assert(moduleOp && "cannot find gpu.module op");
+
+  // Use already generated global op if it exists
+  int index = 0;
+  std::string prefix = llvm::formatv("__shmem_{0}", funcOp.getSymName());
+  moduleOp->walk([&](LLVM::GlobalOp globalOp) {
+    if (auto arrayType = dyn_cast<LLVM::LLVMArrayType>(globalOp.getType())) {
+      if (arrayType.getNumElements() == 0) {
+        existingGlobalOp = globalOp;
+        return WalkResult::interrupt();
+      }
+    }
+    if (globalOp.getSymName().startswith(prefix))
+      index++;
----------------
grypp wrote:

Here what I mean by that. The IR isn't fully lowered yet by the time this PR's patterns run, so I guess that's why I cannot build a symbol table (I get errors above) 
 
```
"gpu.module"() ({
  "llvm.func"() <{CConv = #llvm.cconv<ccc>, function_type = !llvm.func<void (i64)>, linkage = #llvm.linkage<external>, sym_name = "dynamic_shared_memory_kernel", visibility_ = 0 : i64}> ({
  ^bb0(%arg0: i64):
    %0 = "llvm.mlir.constant"() <{value = 8192 : index}> : () -> i64
    %1 = "arith.constant"() <{value = 8192 : index}> : () -> index
    %2 = "llvm.mlir.constant"() <{value = 16384 : index}> : () -> i64
    %3 = "arith.constant"() <{value = 16384 : index}> : () -> index
    %4 = "gpu.dynamic_shared_memory"() : () -> memref<?xi8, #gpu.address_space<workgroup>>
    %5 = "gpu.dynamic_shared_memory"() : () -> memref<?xi8, #gpu.address_space<workgroup>>
    %6 = "memref.view"(%4, %1) : (memref<?xi8, #gpu.address_space<workgroup>>, index) -> memref<32x64xf32, #gpu.address_space<workgroup>>
    "test.use.shared.memory"(%6) : (memref<32x64xf32, #gpu.address_space<workgroup>>) -> ()
    %7 = "memref.view"(%4, %3) : (memref<?xi8, #gpu.address_space<workgroup>>, index) -> memref<32x64xf32, #gpu.address_space<workgroup>>
    "test.use.shared.memory"(%7) : (memref<32x64xf32, #gpu.address_space<workgroup>>) -> ()
    "gpu.return"() : () -> ()
  }) {gpu.kernel, gpu.known_block_size = array<i32: 1, 1, 1>, gpu.known_grid_size = array<i32: 1, 1, 1>, nvvm.kernel} : () -> ()
  "gpu.func"() <{function_type = (index) -> ()}> ({
  }) {gpu.kernel, gpu.known_block_size = array<i32: 1, 1, 1>, gpu.known_grid_size = array<i32: 1, 1, 1>, sym_name = "dynamic_shared_memory_kernel", workgroup_attributions = 0 : i64} : () -> ()
  "gpu.module_end"() : () -> ()
}) {sym_name = "modules"} : () -> ()
```

https://github.com/llvm/llvm-project/pull/71546