[ROCm] Add portable HIP build for ZhiLight on AMD GPUs

3rd/bmengine/CMakeLists.txt

@@ -3,7 +3,11 @@ project(main VERSION 0.1)
 
 enable_language(C)
 enable_language(CXX)
-enable_language(CUDA)
+if(USE_HIP)
+  enable_language(HIP)
+else()
+  enable_language(CUDA)
+endif()
 
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED True)

3rd/bmengine/bmengine/CMakeLists.txt

@@ -3,28 +3,52 @@ project(bmengine VERSION 0.1)
 
 enable_language(C)
 enable_language(CXX)
-if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
-  set(CMAKE_CUDA_ARCHITECTURES "80;89;90a")
-endif()
-if(NOT APPLE)
-  enable_language(CUDA)
+if(USE_HIP)
+  # enable_language(HIP) honors -DCMAKE_HIP_ARCHITECTURES, otherwise auto-detects the
+  # host GPU(s) and errors if none is found; do not pin a default arch here.
+  enable_language(HIP)
+else()
+  if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
+    set(CMAKE_CUDA_ARCHITECTURES "80;89;90a")
+  endif()
+  if(NOT APPLE)
+    enable_language(CUDA)
+  endif()
 endif()
 
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED True)
 set(CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake/)
 
 
-find_library(CUDART_LIBRARY cudart ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES})
-find_library(CUBLAS_LIBRARY cublas ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES})
-find_library(CUBLASLT_LIBRARY cublasLt ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES})
-find_library(CURAND_LIBRARY curand ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES})
-find_library(CULIBOS_LIBRARY culibos ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES})
+if(USE_HIP)
+  # cuBLAS/cuBLASLt/cuRAND -> hipBLAS/hipBLASLt/hipRAND (NCCL -> RCCL below).
+  # culibos has no ROCm analogue (its symbols live in the HIP runtime).
+  find_library(HIPBLAS_LIBRARY hipblas HINTS /opt/rocm/lib)
+  find_library(HIPBLASLT_LIBRARY hipblaslt HINTS /opt/rocm/lib)
+  find_library(HIPRAND_LIBRARY hiprand HINTS /opt/rocm/lib)
+  set(CUDART_LIBRARY "")
+  set(CUBLAS_LIBRARY ${HIPBLAS_LIBRARY})
+  set(CUBLASLT_LIBRARY ${HIPBLASLT_LIBRARY})
+  set(CURAND_LIBRARY ${HIPRAND_LIBRARY})
+  set(CULIBOS_LIBRARY "")
+else()
+  find_library(CUDART_LIBRARY cudart ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES})
+  find_library(CUBLAS_LIBRARY cublas ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES})
+  find_library(CUBLASLT_LIBRARY cublasLt ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES})
+  find_library(CURAND_LIBRARY curand ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES})
+  find_library(CULIBOS_LIBRARY culibos ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES})
+endif()
 
 # ENABLE_NCCL_TP=on by default.
 if(NOT DEFINED ENV{ENABLE_NCCL_TP} OR "$ENV{ENABLE_NCCL_TP}" STREQUAL "on")
     add_definitions(-DENABLE_NCCL_TP=1)
-    find_package(NCCL REQUIRED)
+    if(USE_HIP)
+        find_library(NCCL_LIBRARIES rccl HINTS /opt/rocm/lib REQUIRED)
+        set(NCCL_INCLUDE_DIRS /opt/rocm/include/rccl /opt/rocm/include)
+    else()
+        find_package(NCCL REQUIRED)
+    endif()
     include_directories(SYSTEM ${NCCL_INCLUDE_DIRS})
 else()
     message(STATUS "Not compiling with NCCL support.")
@@ -44,6 +68,12 @@ set_property(TARGET bmengine PROPERTY POSITION_INDEPENDENT_CODE ON)
 set_property(TARGET bmengine PROPERTY CMAKE_CXX_VISIBILITY_PRESET hidden)
 set_property(TARGET bmengine PROPERTY CMAKE_CUDA_VISIBILITY_PRESET hidden)
 
+if(USE_HIP)
+    set_source_files_properties(${FILES_BMENGINE_CUDA} PROPERTIES LANGUAGE HIP)
+    set_target_properties(bmengine PROPERTIES HIP_ARCHITECTURES "${CMAKE_HIP_ARCHITECTURES}")
+    target_compile_options(bmengine PRIVATE $<$<COMPILE_LANGUAGE:HIP>:--offload-compress>)
+endif()
+
 target_include_directories(bmengine 
     PUBLIC "include"
     PUBLIC "include/private/3rd/"

3rd/bmengine/bmengine/core/engine.cpp

@@ -70,13 +70,18 @@ DeviceHandles::DeviceHandles(int dev_id, ncclUniqueId uniqueID, int tp_rank, int
         std::cout << "CC:" << compute_capability
             << ", mp_count:" << mp_count
             << ", L2 Cache:" << (l2_cache_size / 1024 / 1024) << "MB"
+#if !defined(USE_HIP)
             << ", Max Persistent L2:" << (dev_prop.persistingL2CacheMaxSize / 1024) << "KB"
+#endif
             << ", max_smem:" << (max_shared_memory / 1024) << "KB\n";
+#if !defined(USE_HIP)
+    // The persisting-L2-cache device limit is not programmable on CDNA/gfx90a.
     int max_persist_l2 = get_int_env("MAX_PERSIST_L2", 0);
     if (max_persist_l2)  {
         BM_CUDART_ASSERT(
             cudaDeviceSetLimit(cudaLimitPersistingL2CacheSize, max_persist_l2 * 1024 * 1024));
     }
+#endif
 }
 DeviceHandles::~DeviceHandles() {
     DeviceGuard guard(dev_id);

3rd/bmengine/bmengine/core/exception.cpp

@@ -14,11 +14,14 @@ const char* cublasGetErrorString(cublasStatus_t status) {
         case CUBLAS_STATUS_ALLOC_FAILED: return "CUBLAS_STATUS_ALLOC_FAILED";
         case CUBLAS_STATUS_INVALID_VALUE: return "CUBLAS_STATUS_INVALID_VALUE";
         case CUBLAS_STATUS_ARCH_MISMATCH: return "CUBLAS_STATUS_ARCH_MISMATCH";
-        case CUBLAS_STATUS_MAPPING_ERROR: return "CUBLAS_STATUS_MAPPING_ERROR";
         case CUBLAS_STATUS_EXECUTION_FAILED: return "CUBLAS_STATUS_EXECUTION_FAILED";
-        case CUBLAS_STATUS_INTERNAL_ERROR: return "CUBLAS_STATUS_INTERNAL_ERROR";
         case CUBLAS_STATUS_NOT_SUPPORTED: return "CUBLAS_STATUS_NOT_SUPPORTED";
+#if !defined(USE_HIP)
+        // hipBLAS has no MAPPING/INTERNAL/LICENSE status codes.
+        case CUBLAS_STATUS_MAPPING_ERROR: return "CUBLAS_STATUS_MAPPING_ERROR";
+        case CUBLAS_STATUS_INTERNAL_ERROR: return "CUBLAS_STATUS_INTERNAL_ERROR";
         case CUBLAS_STATUS_LICENSE_ERROR: return "CUBLAS_STATUS_LICENSE_ERROR";
+#endif
         default: return "Unknown CUBLAS error";
     }
 }

3rd/bmengine/bmengine/functions/gemm.cpp

@@ -31,9 +31,11 @@ struct LtGemmAlgoAttr {
     int wave_count;
 };
 
+#if !defined(USE_HIP)
 LtGemmAlgoAttr get_algo_attr(const cublasLtMatmulAlgo_t* algo);
 std::string tile_to_str(int id);
 std::string stage_to_str(int id);
+#endif
 
 // clang-format off
 class Gemm::impl {
@@ -71,8 +73,16 @@ class Gemm::impl {
             compute_type = CUBLAS_COMPUTE_32F;
             scale_type = in_type = out_type = CUDA_R_32F;
         } else if (data_type == core::DataType::kHalf) {
+#if defined(USE_HIP)
+            // hipBLASLt on gfx90a rejects 16F compute/scale (INVALID_VALUE);
+            // use 32F compute + 32F scale with fp16 I/O (matches the bf16 path).
+            compute_type = CUBLAS_COMPUTE_32F;
+            scale_type = CUDA_R_32F;
+            in_type = out_type = CUDA_R_16F;
+#else
             compute_type = CUBLAS_COMPUTE_16F;
             scale_type = in_type = out_type = CUDA_R_16F;
+#endif
         } else if (data_type == core::DataType::kBFloat16) {
             // Only support 32F
             compute_type = CUBLAS_COMPUTE_32F;
@@ -173,6 +183,50 @@ class Gemm::impl {
         }
     }
 
+#if defined(USE_HIP)
+    // hipBLASLt has no cublasLtMatmulAlgoInit / AlgoConfigGetAttribute / tile+
+    // stage enums / search-by-algo-id. Correctness-first: ask hipBLASLt for the
+    // best heuristic algo (zero-workspace preference) and cache it. The cuBLASLt
+    // algo-id filtering and tile/stage debug strings are NVIDIA-only and dropped.
+    void find_algo(
+        const core::Context& ctx,
+        cublasLtMatrixLayout_t layout_A,
+        cublasLtMatrixLayout_t layout_B,
+        cublasLtMatrixLayout_t layout_C,
+        uint32_t M,
+        uint32_t K,
+        uint32_t N,
+        uint32_t batch = 0) {
+        if (algo_id != -1 && (last_m != M || last_batch != batch)) {
+            algo_found = false;
+            hipblasLtMatmulPreference_t preference;
+            BM_CUBLAS_ASSERT(hipblasLtMatmulPreferenceCreate(&preference));
+            uint64_t workspace_size = 0;
+            BM_CUBLAS_ASSERT(hipblasLtMatmulPreferenceSetAttribute(
+                preference, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
+                &workspace_size, sizeof(workspace_size)));
+            cublasLtMatmulDesc_t matmul_desc = create_desc();
+            const int kRequest = 4;
+            hipblasLtMatmulHeuristicResult_t results[kRequest];
+            int nb_result = 0;
+            auto status = hipblasLtMatmulAlgoGetHeuristic(
+                ctx.current_cublas_handle(), matmul_desc,
+                layout_B, layout_A, layout_C, layout_C,
+                preference, kRequest, &results[0], &nb_result);
+            if (status == HIPBLAS_STATUS_SUCCESS && nb_result > 0) {
+                algo = results[0].algo;
+                algo_found = true;
+            } else if (last_m == 0) {
+                std::cerr << prefix << ", N=" << N << ", K=" << K << ", M=" << M
+                          << ", batch=" << batch << ". hipBLASLt heuristic found no algo.\n";
+            }
+            hipblasLtMatmulPreferenceDestroy(preference);
+            hipblasLtMatmulDescDestroy(matmul_desc);
+            last_m = M;
+            last_batch = batch;
+        }
+    }
+#else
     void find_algo(
         const core::Context& ctx,
         cublasLtMatrixLayout_t layout_A,
@@ -254,6 +308,7 @@ class Gemm::impl {
             last_batch = batch;
         }
     }
+#endif // USE_HIP
 
     core::Tensor gemm(
         const core::Context& ctx,
@@ -599,6 +654,7 @@ void Gemm::set_B_scale(const core::Tensor& B_scale) {
     pimpl->b_scale = B_scale.data();
 }
 
+#if !defined(USE_HIP)
 LtGemmAlgoAttr get_algo_attr(const cublasLtMatmulAlgo_t* algo) {
     LtGemmAlgoAttr attr;
     BM_CUBLAS_ASSERT(cublasLtMatmulAlgoConfigGetAttribute(
@@ -683,6 +739,7 @@ std::string stage_to_str(int id) {
     };
     return stage_map.at(static_cast<cublasLtMatmulStages_t>(id)).substr(23);
 }
+#endif // !USE_HIP
 // clang-format on
 
 } // namespace functions

3rd/bmengine/bmengine/functions/sort.cu

@@ -11,7 +11,34 @@ namespace functions {
 using bmengine::core::Tensor;
 
 // https://nvidia.github.io/cccl/cub/api/structcub_1_1DeviceRadixSort.html
+//
+// rocPRIM's DeviceRadixSort (hipCUB backend) only accepts native radix-sortable
+// key types; half/__hip_bfloat16 keys require a custom decomposer (CUB takes
+// __half directly). The half/bf16 key sort is not on the inference fast path,
+// so on HIP it is left unsupported (runtime error if ever reached) rather than
+// shipping an order-incorrect bit-cast sort.
 template<typename KeyT>
+struct RadixSortable { static constexpr bool value = true; };
+#if defined(USE_HIP)
+template<> struct RadixSortable<half> { static constexpr bool value = false; };
+template<> struct RadixSortable<nv_bfloat16> { static constexpr bool value = false; };
+#endif
+
+template<typename KeyT,
+         typename std::enable_if<!RadixSortable<KeyT>::value, int>::type = 0>
+std::pair<core::Tensor, core::Tensor> sort_1d_temp(
+    const core::Context& ctx,
+    const core::Tensor& keys,
+    const core::Tensor& values,
+    int max_key)
+{
+    BM_EXCEPTION("sort_pair_1d with half/bfloat16 keys is not supported on ROCm/HIP "
+                 "(rocPRIM radix sort needs a decomposer for these key types).");
+    return {};
+}
+
+template<typename KeyT,
+         typename std::enable_if<RadixSortable<KeyT>::value, int>::type = 0>
 std::pair<core::Tensor, core::Tensor> sort_1d_temp(
     const core::Context& ctx,
     const core::Tensor& keys,

3rd/bmengine/bmengine/functions/topk.cu

@@ -5,6 +5,19 @@ namespace bmengine {
 
 namespace functions {
 
+// XOR-permute within the warp/wavefront. On HIP the no-mask form spans the full
+// (32- or 64-lane) wavefront; CUDA keeps the full-warp 32-bit mask. The bitonic
+// groups are width-N (N <= wave size), so lane^j stays inside the wavefront and
+// the permute is correct on wave32 and wave64.
+template<typename T>
+static __device__ inline T warpXorShfl(T v, int laneMask) {
+#if defined(__HIP_PLATFORM_AMD__) || defined(USE_HIP)
+    return __shfl_xor(v, laneMask);
+#else
+    return __shfl_xor_sync(0xFFFFFFFF, v, laneMask);
+#endif
+}
+
 template<typename T, int N>
 static __device__ inline void warpBitonicSort(T& v1, int& pos, bool asc) {
     int lane_id = threadIdx.x & (N - 1);
@@ -13,8 +26,8 @@ static __device__ inline void warpBitonicSort(T& v1, int& pos, bool asc) {
         bool desc = ((lane_id & k) == 0) ^ asc;
 #pragma unroll
         for (int j = k / 2; j > 0; j /= 2) {
-            T v2 = __shfl_xor_sync(0xFFFFFFFF, v1, j);
-            int pos2 = __shfl_xor_sync(0xFFFFFFFF, pos, j);
+            T v2 = warpXorShfl(v1, j);
+            int pos2 = warpXorShfl(pos, j);
             bool upper = (lane_id & j) != 0;
 
             if (desc ^ (v1 > v2 || (v1 == v2 && pos < pos2)) ^ upper) {
@@ -35,8 +48,8 @@ static __device__ inline void warpBitonicMerge(T& v1, int& pos1, T& v2, int& pos
 // resort
 #pragma unroll
     for (int j = N / 2; j > 0; j /= 2) {
-        v2 = __shfl_xor_sync(0xFFFFFFFF, v1, j);
-        int pos2 = __shfl_xor_sync(0xFFFFFFFF, pos1, j);
+        v2 = warpXorShfl(v1, j);
+        int pos2 = warpXorShfl(pos1, j);
         bool upper = (lane_id & j) != 0;
         if ((v1 < v2 || (v1 == v2 && pos1 > pos2)) ^ upper) {
             v1 = v2;

3rd/bmengine/bmengine/include/bmengine/core/utils.h

@@ -48,7 +48,15 @@ inline T ceil_div(T m, Tb d) {
 
 template<typename T>
 inline T round_up_thread(T m) {
+    // Round the block size up to a whole wavefront/warp. On CDNA (gfx90a) the
+    // wavefront is 64, so rounding to 32 would launch a half-filled wavefront;
+    // the full-wavefront shuffles in the warp/block reductions then read
+    // inactive lanes and give wrong amax/amin/etc. Round to 64 on HIP.
+#if defined(USE_HIP)
+    T d = 64, limit = MAX_NUM_THREADS;
+#else
     T d = 32, limit = MAX_NUM_THREADS;
+#endif
     T x = m > limit ? limit : m;
     return ((x + d - 1) / d) * d;
 }