[UR][HIP] Opt into the larger per-block shared-memory pool

sycl/test-e2e/Adapters/hip-local-mem-optin.cpp

@@ -0,0 +1,180 @@
+// REQUIRES: hip, hip_dev_kit
+// UNSUPPORTED: windows
+// UNSUPPORTED-INTENDED: %{build} targets spir64, so --rocm-path from
+// %hip_options is an unused argument and fails under -Werror.
+
+// RUN: %{build} -o %t.out %hip_options
+// The opt-in launch is skipped when the opt-in ceiling equals the non-opt-in
+// limit. Otherwise it uses one byte more than that limit.
+// RUN: env SYCL_DEVICE_FILTER=hip %{run} %t.out
+// A cap below the device ceiling lowers the query and rejects a larger launch.
+// RUN: env SYCL_DEVICE_FILTER=hip UR_HIP_MAX_LOCAL_MEM_SIZE=32768 %{run} %t.out
+
+//==------------------------ hip-local-mem-optin.cpp ----------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+// Dynamic local memory above the non-opt-in per-block limit has to opt in with
+// hipFuncAttributeMaxDynamicSharedMemorySize. AMD devices report the same
+// ceiling for both attributes (64KiB LDS on gfx942), so that launch is skipped
+// there. The test still checks the reported ceiling, a launch of that size,
+// and rejection of one byte more. UR_HIP_MAX_LOCAL_MEM_SIZE caps both.
+
+#include <sycl/detail/core.hpp>
+#include <sycl/ext/intel/info/device.hpp>
+
+#include <algorithm>
+#include <cstdio>
+#include <cstdlib>
+#include <string>
+
+#if !defined(__HIP_PLATFORM_AMD__) && !defined(__HIP_PLATFORM_NVIDIA__)
+#define __HIP_PLATFORM_AMD__
+#endif
+#include <hip/hip_runtime_api.h>
+
+namespace {
+
+int queryAttribute(int Device, hipDeviceAttribute_t Attr, bool *Ok) {
+  int Value = 0;
+  hipError_t Err = hipDeviceGetAttribute(&Value, Attr, Device);
+  if (Err != hipSuccess) {
+    (void)hipGetLastError();
+    *Ok = false;
+    return 0;
+  }
+  *Ok = true;
+  return Value;
+}
+
+bool launchAndCheck(sycl::queue &Q, size_t Bytes) {
+  sycl::buffer<unsigned char, 1> Out{sycl::range<1>(2)};
+  Q.submit([&](sycl::handler &CGH) {
+     sycl::accessor OutAcc{Out, CGH, sycl::write_only};
+     sycl::local_accessor<unsigned char, 1> Local{Bytes, CGH};
+     CGH.parallel_for(sycl::nd_range<1>{1, 1}, [=](sycl::nd_item<1>) {
+       Local[0] = 0xA5;
+       Local[Bytes - 1] = 0x5A;
+       OutAcc[0] = Local[0];
+       OutAcc[1] = Local[Bytes - 1];
+     });
+   }).wait();
+  sycl::host_accessor Host{Out, sycl::read_only};
+  if (Host[0] != 0xA5 || Host[1] != 0x5A) {
+    std::printf("launch %zu read back %u %u\n", Bytes, Host[0], Host[1]);
+    return false;
+  }
+  std::printf("launch %zu OK\n", Bytes);
+  return true;
+}
+
+bool expectReject(sycl::queue &Q, size_t Bytes, const char *Expected) {
+  try {
+    launchAndCheck(Q, Bytes);
+  } catch (const std::exception &E) {
+    if (std::string(E.what()).find(Expected) == std::string::npos) {
+      std::printf("launch %zu rejected with unexpected message: %s\n", Bytes,
+                  E.what());
+      return false;
+    }
+    std::printf("launch %zu rejected as expected\n", Bytes);
+    return true;
+  }
+  std::printf("launch %zu was expected to fail\n", Bytes);
+  return false;
+}
+
+} // namespace
+
+int main() {
+  sycl::queue Q;
+  const std::string Pci =
+      Q.get_device().get_info<sycl::ext::intel::info::device::pci_address>();
+  int DeviceCount = 0;
+  if (hipGetDeviceCount(&DeviceCount) != hipSuccess || DeviceCount <= 0) {
+    std::printf("failed to enumerate HIP devices\n");
+    return 1;
+  }
+  int HipDevice = -1;
+  for (int Index = 0; Index < DeviceCount; ++Index) {
+    char BusId[32] = {};
+    if (hipDeviceGetPCIBusId(BusId, sizeof(BusId), Index) == hipSuccess &&
+        Pci == BusId) {
+      HipDevice = Index;
+      break;
+    }
+  }
+  if (HipDevice < 0) {
+    std::printf("no HIP device matches PCI address %s\n", Pci.c_str());
+    return 1;
+  }
+
+  bool DefaultOk = false;
+  const int Default = queryAttribute(
+      HipDevice, hipDeviceAttributeMaxSharedMemoryPerBlock, &DefaultOk);
+  if (!DefaultOk || Default <= 0) {
+    std::printf("failed to query non-opt-in shared memory on HIP device %d\n",
+                HipDevice);
+    return 1;
+  }
+
+  bool OptinOk = false;
+  int Optin = queryAttribute(
+      HipDevice, hipDeviceAttributeSharedMemPerBlockOptin, &OptinOk);
+  // Match the adapter: a rejected or non-positive opt-in query keeps the
+  // non-opt-in size.
+  if (!OptinOk || Optin <= 0)
+    Optin = Default;
+  const int Ceiling = std::max(Default, Optin);
+
+  const auto Reported = static_cast<size_t>(
+      Q.get_device().get_info<sycl::info::device::local_mem_size>());
+
+  const char *CapEnv = std::getenv("UR_HIP_MAX_LOCAL_MEM_SIZE");
+  size_t Expected = static_cast<size_t>(Ceiling);
+  if (CapEnv) {
+    const int Cap = std::atoi(CapEnv);
+    if (Cap <= 0) {
+      std::printf("UR_HIP_MAX_LOCAL_MEM_SIZE is not a positive size\n");
+      return 1;
+    }
+    Expected = static_cast<size_t>(std::min(Cap, Ceiling));
+  }
+
+  std::printf("device %d default %d opt-in %d ceiling %d reported %zu\n",
+              HipDevice, Default, Optin, Ceiling, Reported);
+  if (Reported != Expected) {
+    std::printf("local_mem_size %zu, expected %zu\n", Reported, Expected);
+    return 1;
+  }
+
+  if (Reported == 0 || !launchAndCheck(Q, Reported))
+    return 1;
+
+  const char *Reject =
+      (Reported + 1 > static_cast<size_t>(Ceiling))
+          ? "Excessive allocation of local memory on the device"
+          : "UR_HIP_MAX_LOCAL_MEM_SIZE";
+  if (!expectReject(Q, Reported + 1, Reject))
+    return 1;
+
+  // One byte past the non-opt-in limit is the hipFuncSetAttribute path. It is
+  // unreachable when the ceiling equals that limit (gfx942 and the other AMD
+  // ISAs) or when the env cap keeps the launch inside the default pool.
+  const size_t OptInBytes = static_cast<size_t>(Default) + 1;
+  if (Ceiling > Default && OptInBytes <= Reported && OptInBytes != Reported) {
+    if (!launchAndCheck(Q, OptInBytes))
+      return 1;
+    std::printf("opt-in launch %zu OK\n", OptInBytes);
+  } else if (Ceiling > Default && OptInBytes == Reported) {
+    std::printf("opt-in launch covered by the %zu-byte launch\n", Reported);
+  } else {
+    std::printf("skipped opt-in launch: ceiling %d default %d reported %zu\n",
+                Ceiling, Default, Reported);
+  }
+  return 0;
+}

unified-runtime/source/adapters/hip/device.cpp

@@ -402,7 +402,9 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
   case UR_DEVICE_INFO_LOCAL_MEM_SIZE: {
     // OpenCL's "local memory" maps most closely to HIP's "shared memory".
     // HIP has its own definition of "local memory", which maps to OpenCL's
-    // "private memory".
+    // "private memory". The size is the opt-in per-block ceiling
+    // (hipDeviceAttributeSharedMemPerBlockOptin, or the non-opt-in size when
+    // that attribute is absent). UR_HIP_MAX_LOCAL_MEM_SIZE caps it.
     if (hDevice->getMaxChosenLocalMem()) {
       return ReturnValue(
           static_cast<uint64_t>(hDevice->getMaxChosenLocalMem()));

unified-runtime/source/adapters/hip/device.hpp

@@ -13,6 +13,8 @@
 
 #include <ur/ur.hpp>
 
+#include <algorithm>
+
 /// UR device mapping to a hipDevice_t.
 /// Includes an observer pointer to the platform,
 /// and implements the reference counting semantics since
@@ -28,7 +30,11 @@ struct ur_device_handle_t_ : ur::hip::handle_base {
 
   int MaxWorkGroupSize{0};
   size_t MaxBlockDim[3];
+  // Opt-in per-block shared memory ceiling reported as local_mem_size.
   int MaxCapacityLocalMem{0};
+  // Non-opt-in per-block shared memory. Dynamic local memory above this
+  // requires hipFuncAttributeMaxDynamicSharedMemorySize.
+  int MaxDefaultLocalMem{0};
   int MaxChosenLocalMem{0};
   int ManagedMemSupport{0};
   int ConcurrentManagedAccess{0};
@@ -54,9 +60,31 @@ struct ur_device_handle_t_ : ur::hip::handle_base {
         &MaxDim, hipDeviceAttributeMaxBlockDimZ, HIPDevice));
     MaxBlockDim[2] = size_t(MaxDim);
 
+    int MaxSharedMemPerBlock = 0;
     UR_CHECK_ERROR(hipDeviceGetAttribute(
-        &MaxCapacityLocalMem, hipDeviceAttributeMaxSharedMemoryPerBlock,
+        &MaxSharedMemPerBlock, hipDeviceAttributeMaxSharedMemoryPerBlock,
         HIPDevice));
+    // Default dynamic shared-memory limit. Kernels that need more must opt in
+    // via hipFuncAttributeMaxDynamicSharedMemorySize, up to the opt-in ceiling.
+    MaxDefaultLocalMem = MaxSharedMemPerBlock;
+
+    int MaxSharedMemOptin = 0;
+    hipError_t OptinStatus = hipDeviceGetAttribute(
+        &MaxSharedMemOptin, hipDeviceAttributeSharedMemPerBlockOptin,
+        HIPDevice);
+    if (OptinStatus != hipSuccess) {
+      // ROCm 5.7 rejects this attribute on AMD (hipErrorInvalidValue). Clear
+      // HIP's last-error state so it does not contaminate a later error check.
+      (void)hipGetLastError();
+      MaxSharedMemOptin = 0;
+    }
+    // ROCm 6.2 accepts the query on gfx90a but returns 0. CDNA parts that do
+    // report an opt-in value (ROCm 6.4+) report the same 64KiB LDS size as
+    // the non-opt-in limit, not a larger pool. Keep the larger of the two so
+    // a missing or zero opt-in cannot hide memory the device already exposes.
+    if (MaxSharedMemOptin <= 0)
+      MaxSharedMemOptin = MaxSharedMemPerBlock;
+    MaxCapacityLocalMem = std::max(MaxSharedMemPerBlock, MaxSharedMemOptin);
     UR_CHECK_ERROR(hipDeviceGetAttribute(
         &ManagedMemSupport, hipDeviceAttributeManagedMemory, HIPDevice));
     UR_CHECK_ERROR(hipDeviceGetAttribute(
@@ -114,6 +142,8 @@ struct ur_device_handle_t_ : ur::hip::handle_base {
 
   int getMaxCapacityLocalMem() const noexcept { return MaxCapacityLocalMem; };
 
+  int getMaxDefaultLocalMem() const noexcept { return MaxDefaultLocalMem; };
+
   int getMaxChosenLocalMem() const noexcept { return MaxChosenLocalMem; };
 
   int getManagedMemSupport() const noexcept { return ManagedMemSupport; };

unified-runtime/source/adapters/hip/enqueue.cpp

@@ -1858,27 +1858,8 @@ setKernelParams(const ur_device_handle_t Device, const uint32_t WorkDim,
       Kernel->setImplicitOffsetArg(sizeof(ImplicitOffset), ImplicitOffset);
     }
 
-    uint32_t LocalSize = Kernel->getLocalSize();
-    if (LocalSize > static_cast<uint32_t>(Device->getMaxCapacityLocalMem())) {
-      setErrorMessage("Excessive allocation of local memory on the device",
-                      UR_RESULT_ERROR_OUT_OF_RESOURCES);
-      return UR_RESULT_ERROR_ADAPTER_SPECIFIC;
-    }
-
-    if (int MaxLocalMem = Device->getMaxChosenLocalMem()) {
-      if (LocalSize > static_cast<uint32_t>(MaxLocalMem)) {
-        setErrorMessage(
-            "Local memory for kernel exceeds the amount requested using "
-            "UR_HIP_MAX_LOCAL_MEM_SIZE (or deprecated "
-            "SYCL_PI_HIP_MAX_LOCAL_MEM_SIZE). Try increasing the maximum "
-            "local memory.",
-            UR_RESULT_ERROR_OUT_OF_RESOURCES);
-        return UR_RESULT_ERROR_ADAPTER_SPECIFIC;
-      }
-
-      UR_CHECK_ERROR(hipFuncSetAttribute(
-          HIPFunc, hipFuncAttributeMaxDynamicSharedMemorySize, MaxLocalMem));
-    }
+    Kernel->getProgram()->prepareFunctionLocalMemory(HIPFunc,
+                                                     Kernel->getLocalSize());
   } catch (ur_result_t Err) {
     return Err;
   }

unified-runtime/source/adapters/hip/program.cpp

@@ -10,6 +10,58 @@
 #include "program.hpp"
 #include "ur_util.hpp"
 
+void ur_program_handle_t_::prepareFunctionLocalMemory(hipFunction_t Function,
+                                                      uint32_t DynamicSize) {
+  std::lock_guard<std::mutex> Lock(FunctionLocalMemMutex);
+  auto &Info = FunctionLocalMem[Function];
+
+  if (Info.StaticSize < 0) {
+    UR_CHECK_ERROR(hipFuncGetAttribute(
+        &Info.StaticSize, HIP_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES, Function));
+  }
+
+  const uint64_t TotalSize =
+      static_cast<uint64_t>(Info.StaticSize) + DynamicSize;
+  const int Capacity = Device->getMaxCapacityLocalMem();
+  if (TotalSize > static_cast<uint64_t>(Capacity)) {
+    setErrorMessage("Excessive allocation of local memory on the device",
+                    UR_RESULT_ERROR_OUT_OF_RESOURCES);
+    throw UR_RESULT_ERROR_ADAPTER_SPECIFIC;
+  }
+
+  const int ChosenSize = Device->getMaxChosenLocalMem();
+  if (ChosenSize && DynamicSize > static_cast<uint32_t>(ChosenSize)) {
+    setErrorMessage(
+        "Local memory for kernel exceeds the amount requested using "
+        "UR_HIP_MAX_LOCAL_MEM_SIZE (or deprecated "
+        "SYCL_PI_HIP_MAX_LOCAL_MEM_SIZE). Try increasing the maximum "
+        "local memory.",
+        UR_RESULT_ERROR_OUT_OF_RESOURCES);
+    throw UR_RESULT_ERROR_ADAPTER_SPECIFIC;
+  }
+
+  if (TotalSize <= static_cast<uint64_t>(Device->getMaxDefaultLocalMem()) ||
+      Info.MaxDynamicSizeSet)
+    return;
+
+  // Set one stable limit for this module function. Multiple UR kernel handles
+  // can share Function, so changing the limit to each launch's requested size
+  // would race with launches through the other handles.
+  const int MaxDynamicSize =
+      std::min(ChosenSize ? ChosenSize : Capacity, Capacity - Info.StaticSize);
+#if defined(__HIP_PLATFORM_NVIDIA__)
+  // hipModuleGetFunction returns a driver function. HIP's NVIDIA runtime
+  // hipFuncSetAttribute wrapper expects a runtime-registered kernel symbol.
+  UR_CHECK_ERROR(hipCUResultTohipError(cuFuncSetAttribute(
+      Function, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
+      MaxDynamicSize)));
+#else
+  UR_CHECK_ERROR(hipFuncSetAttribute(
+      Function, hipFuncAttributeMaxDynamicSharedMemorySize, MaxDynamicSize));
+#endif
+  Info.MaxDynamicSizeSet = true;
+}
+
 #ifdef SYCL_ENABLE_KERNEL_FUSION
 #ifdef UR_COMGR_VERSION4_INCLUDE
 #include <amd_comgr.h>

unified-runtime/source/adapters/hip/program.hpp

@@ -11,6 +11,7 @@
 #include <unified-runtime/ur_api.h>
 
 #include <atomic>
+#include <mutex>
 #include <unordered_map>
 
 #include "common/ur_ref_count.hpp"
@@ -41,6 +42,13 @@ struct ur_program_handle_t_ : ur::hip::handle_base {
       KernelReqdWorkGroupSizeMD;
   std::unordered_map<std::string, uint32_t> KernelReqdSubGroupSizeMD;
 
+  struct FunctionLocalMemInfo {
+    int StaticSize{-1};
+    bool MaxDynamicSizeSet{false};
+  };
+  std::mutex FunctionLocalMemMutex;
+  std::unordered_map<hipFunction_t, FunctionLocalMemInfo> FunctionLocalMem;
+
   constexpr static size_t MAX_LOG_SIZE = 8192u;
 
   char ErrorLog[MAX_LOG_SIZE], InfoLog[MAX_LOG_SIZE];
@@ -66,6 +74,7 @@ struct ur_program_handle_t_ : ur::hip::handle_base {
 
   ur_result_t buildProgram(const char *BuildOptions);
   ur_result_t finalizeRelocatable();
+  void prepareFunctionLocalMemory(hipFunction_t Function, uint32_t DynamicSize);
   ur_context_handle_t getContext() const { return Context; };
   ur_device_handle_t getDevice() const { return Device; };