[Codegen][ROCDL] Drop nominal support for dynamic shared mem (#20805)

This was introduced for CUDA to work around static allocation limits
(see #8317) and then subsequently cloned for ROCM. There is no such
allocation limit on the AMD side and we don't actually support dynamic
shared memory, so we can drop it.

Immediately this fixes issues where multiple exported functions will
request the sum of the memory required across all functions because the
dynamic symbol obfuscates the actual required amount.

To avoid regressions, this matches the current behavior of forcing all
shared memory allocations to be aligned by 16. Runtime
support for dynamic shared is dropped and the `block_shared_memory_size`
field in the flatbuffer is marked as deprecated. To
retain load time checking of shared memory resource constraints, we now
query exported functions directly using `hipFuncGetAttribute(...,
HIP_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES)`.

Author: qedawkins

compiler/plugins/target/ROCM/ROCMTarget.cpp

@@ -828,12 +828,6 @@ class ROCMTargetBackend final : public TargetBackend {
         blockDims.z = cast<IntegerAttr>(workgroupSize[2]).getInt();
       }
 
-      uint32_t blockSharedMemorySize = 0;
-      if (std::optional<APInt> workgroupLocalMemoryAttr =
-              exportOp.getWorkgroupLocalMemory()) {
-        blockSharedMemorySize = workgroupLocalMemoryAttr->getSExtValue();
-      }
-
       auto layoutAttr = exportOp.getLayoutAttr();
       uint32_t constantCount = static_cast<uint32_t>(layoutAttr.getConstants());
       SmallVector<iree_hal_hip_BindingBits_enum_t> bindingFlags;
@@ -856,8 +850,6 @@ class ROCMTargetBackend final : public TargetBackend {
       iree_hal_hip_ExportDef_module_ordinal_add(builder, 0); // always 0 today
       iree_hal_hip_ExportDef_kernel_name_add(builder, kernelNameRef);
       iree_hal_hip_ExportDef_block_dims_add(builder, &blockDims);
-      iree_hal_hip_ExportDef_block_shared_memory_size_add(
-          builder, blockSharedMemorySize);
       iree_hal_hip_ExportDef_constant_count_add(builder, constantCount);
       iree_hal_hip_ExportDef_binding_flags_add(builder, bindingFlagsRef);
       iree_hal_hip_ExportDef_debug_info_add(builder, exportDebugInfos[ordinal]);

compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp

@@ -32,7 +32,7 @@ namespace mlir::iree_compiler {
 
 void ConvertToDynamicSharedMemory(ModuleOp moduleOp) {
   SymbolTableCollection symbolTableCollection;
-  // Collect all the adressOfOps to static shared memory globals.
+  // Collect all the addressOfOps to static shared memory globals.
   SmallVector<LLVM::AddressOfOp> addressOfOps;
   moduleOp.walk([&](LLVM::AddressOfOp addressOfOp) {
     // Check that the global associated with this addressOfOp has shared memory
@@ -95,6 +95,18 @@ void ConvertToDynamicSharedMemory(ModuleOp moduleOp) {
   }
 }
 
+void setSharedMemoryAlignment(ModuleOp moduleOp, uint64_t newAlignment) {
+  for (auto global : moduleOp.getOps<LLVM::GlobalOp>()) {
+    if (global.getAddrSpace() == 3) {
+      uint64_t baseAlignment = 0;
+      if (std::optional<uint64_t> alignment = global.getAlignment()) {
+        baseAlignment = alignment.value();
+      }
+      global.setAlignment(std::max<uint64_t>(baseAlignment, newAlignment));
+    }
+  }
+}
+
 namespace {
 
 /// Scalarize math ops. It is needed to lower vector operation that don't have

compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.h

@@ -29,6 +29,8 @@ void populateConvertSharedMemoryAllocOps(RewritePatternSet &patterns);
 
 void ConvertToDynamicSharedMemory(ModuleOp moduleOp);
 
+void setSharedMemoryAlignment(ModuleOp moduleOp, uint64_t newAlignment);
+
 using MemorySpaceMapping =
     std::function<unsigned(gpu::AddressSpace gpuAddressSpace)>;
 void populateGpuMemorySpaceAttributeConversions(

compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToROCDL.cpp

@@ -339,9 +339,14 @@ struct ConvertToROCDLPass final
     }
 
     LDBG("After converting to rocdl\n" << m);
-    ConvertToDynamicSharedMemory(m);
 
-    LDBG("After converting to dynamic shared memory\n" << m);
+    // 16 is the maximum relevant alignment for all AMD GPUs. Unceremoniously
+    // set it to 16 as all of our allocations almost always have much greater
+    // alignment than this.
+    // TODO(qedawkins): Set this much earlier when we introduce the allocations.
+    setSharedMemoryAlignment(m, 16);
+
+    LDBG("After updating shared memory alignments\n" << m);
   }
 };
 } // namespace mlir::iree_compiler

compiler/src/iree/compiler/Codegen/LLVMGPU/test/convert_to_rocdl.mlir

@@ -305,3 +305,30 @@ module {
 //   CHECK-NEXT:    rocdl.s.setprio 3
 //   CHECK:         rocdl.mfma
 //   CHECK-NEXT:    rocdl.s.setprio 4
+
+// -----
+
+builtin.module {
+  func.func @shared_memory_lowering() {
+    %c0 = arith.constant 0 : index
+    %cst = arith.constant dense<0.000000e+00> : vector<4xf32>
+    %0 = memref.alloc() : memref<1x16x32xf32, #gpu.address_space<workgroup>>
+    %1 = memref.alloc() : memref<1x32x16xf32, #gpu.address_space<workgroup>>
+    %2 = memref.alloc() : memref<1x8x16xf32, #gpu.address_space<workgroup>>
+    vector.store %cst, %1[%c0, %c0, %c0] : memref<1x32x16xf32, #gpu.address_space<workgroup>>, vector<4xf32>
+    vector.store %cst, %2[%c0, %c0, %c0] : memref<1x8x16xf32, #gpu.address_space<workgroup>>, vector<4xf32>
+    vector.store %cst, %0[%c0, %c0, %c0] : memref<1x16x32xf32, #gpu.address_space<workgroup>>, vector<4xf32>
+    return
+  }
+}
+
+// CHECK-DAG: llvm.mlir.global private @__shared_memory___1() {addr_space = 3 : i32, alignment = 16 : i64} : !llvm.array<1 x array<16 x array<32 x f32>>>
+// CHECK-DAG: llvm.mlir.global private @__shared_memory___0() {addr_space = 3 : i32, alignment = 16 : i64} : !llvm.array<1 x array<32 x array<16 x f32>>>
+// CHECK-DAG: llvm.mlir.global private @__shared_memory__() {addr_space = 3 : i32, alignment = 16 : i64} : !llvm.array<1 x array<8 x array<16 x f32>>>
+// CHECK-LABEL: llvm.func @shared_memory_lowering() {
+//   CHECK-DAG: %[[A1:.+]] = llvm.mlir.addressof @__shared_memory___1
+//   CHECK-DAG: llvm.getelementptr %[[A1]][0, 0, 0, 0]
+//   CHECK-DAG: %[[A0:.+]] = llvm.mlir.addressof @__shared_memory___0
+//   CHECK-DAG: llvm.getelementptr %[[A0]][0, 0, 0, 0]
+//   CHECK-DAG: %[[A:.+]] = llvm.mlir.addressof @__shared_memory__
+//   CHECK-DAG: llvm.getelementptr %[[A]][0, 0, 0, 0]

runtime/src/iree/hal/drivers/hip/dynamic_symbol_tables.h

@@ -36,6 +36,8 @@ IREE_HAL_HIP_REQUIRED_PFN_DECL(hipEventRecord, hipEvent_t, hipStream_t)
 IREE_HAL_HIP_REQUIRED_PFN_DECL(hipEventSynchronize, hipEvent_t)
 IREE_HAL_HIP_REQUIRED_PFN_DECL(hipFree, void *)
 IREE_HAL_HIP_REQUIRED_PFN_DECL(hipFreeAsync, void *, hipStream_t)
+IREE_HAL_HIP_REQUIRED_PFN_DECL(hipFuncGetAttribute, int *, hipFuncAttribute,
+                               const void *)
 IREE_HAL_HIP_REQUIRED_PFN_DECL(hipFuncSetAttribute, const void *,
                                hipFuncAttribute, int)
 IREE_HAL_HIP_REQUIRED_PFN_DECL(hipGetDeviceCount, int *)

runtime/src/iree/hal/drivers/hip/graph_command_buffer.c

@@ -815,7 +815,6 @@ static iree_status_t iree_hal_hip_graph_command_buffer_dispatch(
       .gridDim.z = workgroup_count[2],
       .func = kernel_params->function,
       .kernelParams = params_ptr,
-      .sharedMemBytes = kernel_params->block_shared_memory_size,
   };
 
   if (command_buffer->graph_node_count >=

runtime/src/iree/hal/drivers/hip/native_executable.c

@@ -174,17 +174,6 @@ static iree_status_t iree_hal_hip_native_executable_flatbuffer_verify(
                               i);
     }
 
-    uint32_t block_shared_memory_size =
-        iree_hal_hip_ExportDef_block_shared_memory_size_get(export_def);
-    if (block_shared_memory_size > limits->max_block_shared_memory_size) {
-      return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
-                              "exports[%" PRIhsz
-                              "] requires %uB of shared memory and "
-                              "exceeds the device maximum of %uB per block",
-                              i, block_shared_memory_size,
-                              limits->max_block_shared_memory_size);
-    }
-
     uint32_t constant_count =
         iree_hal_hip_ExportDef_constant_count_get(export_def);
     if (constant_count > IREE_HAL_HIP_MAX_DISPATCH_CONSTANT_COUNT) {
@@ -212,6 +201,30 @@ static iree_status_t iree_hal_hip_native_executable_flatbuffer_verify(
   return iree_ok_status();
 }
 
+// Verifies a function against the device limits so that we can avoid doing so
+// during runtime.
+static iree_status_t iree_hal_hip_function_attributes_verify(
+    iree_host_size_t id, const iree_hal_hip_dynamic_symbols_t* symbols,
+    hipFunction_t function, const iree_hal_hip_limits_t* limits) {
+  int block_shared_memory_size;
+  IREE_RETURN_IF_ERROR(IREE_HIP_CALL_TO_STATUS(
+      symbols,
+      hipFuncGetAttribute(
+          &block_shared_memory_size,
+          (hipFuncAttribute)HIP_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES, function),
+      "hipFuncGetAttribute"));
+  if (block_shared_memory_size > limits->max_block_shared_memory_size) {
+    return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
+                            "exports[%" PRIhsz
+                            "] requires %uB of shared memory and "
+                            "exceeds the device maximum of %uB per block",
+                            id, block_shared_memory_size,
+                            limits->max_block_shared_memory_size);
+  }
+
+  return iree_ok_status();
+}
+
 iree_status_t iree_hal_hip_native_executable_create(
     const iree_hal_hip_dynamic_symbols_t* symbols,
     iree_hal_hip_device_topology_t topology,
@@ -397,16 +410,8 @@ iree_status_t iree_hal_hip_native_executable_create(
           break;
         }
 
-        uint32_t block_shared_memory_size =
-            iree_hal_hip_ExportDef_block_shared_memory_size_get(export_def);
-        status = IREE_HIP_CALL_TO_STATUS(
-            symbols,
-            hipFuncSetAttribute(
-                function,
-                (hipFuncAttribute)
-                    HIP_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
-                block_shared_memory_size),
-            "hipFuncSetAttribute");
+        status = iree_hal_hip_function_attributes_verify(i, symbols, function,
+                                                         &limits);
         if (!iree_status_is_ok(status)) break;
 
         // Package required parameters for kernel launches for each entry
@@ -419,8 +424,6 @@ iree_status_t iree_hal_hip_native_executable_create(
         kernel_info->block_dims[0] = block_dims->x;
         kernel_info->block_dims[1] = block_dims->y;
         kernel_info->block_dims[2] = block_dims->z;
-        kernel_info->block_shared_memory_size =
-            iree_hal_hip_ExportDef_block_shared_memory_size_get(export_def);
         kernel_info->constant_count =
             iree_hal_hip_ExportDef_constant_count_get(export_def);
         iree_hal_hip_BindingBits_vec_t binding_flags_vec =

runtime/src/iree/hal/drivers/hip/native_executable.h

@@ -37,7 +37,6 @@ typedef struct iree_hal_hip_kernel_params_t {
   uint32_t binding_count;
 
   uint32_t block_dims[3];
-  uint32_t block_shared_memory_size;
 
   IREE_TRACE(iree_hal_hip_kernel_debug_info_t debug_info;)
 } iree_hal_hip_kernel_params_t;

runtime/src/iree/hal/drivers/hip/stream_command_buffer.c

@@ -578,8 +578,7 @@ static iree_status_t iree_hal_hip_stream_command_buffer_dispatch(
           kernel_params->function, workgroup_count[0], workgroup_count[1],
           workgroup_count[2], kernel_params->block_dims[0],
           kernel_params->block_dims[1], kernel_params->block_dims[2],
-          kernel_params->block_shared_memory_size, command_buffer->hip_stream,
-          params_ptr, NULL),
+          /*sharedMemBytes=*/0, command_buffer->hip_stream, params_ptr, NULL),
       "hipModuleLaunchKernel");
 
   IREE_HAL_STREAM_TRACE_ZONE_END(command_buffer->tracing_context,