nix: Added and enabled rocmWmma support along with ROWMMA_PATH option in CMake

Author: SteelPh0enix

.devops/nix/package.nix

@@ -32,6 +32,7 @@
   useMpi ? false,
   useRocm ? config.rocmSupport,
   rocmGpuTargets ? builtins.concatStringsSep ";" rocmPackages.clr.gpuTargets,
+  rocmUseWmma ? true,
   enableCurl ? true,
   useVulkan ? false,
   buildAllCudaFaQuants ? false,
@@ -92,13 +93,16 @@ let
     libcublas
   ];
 
-  rocmBuildInputs = with rocmPackages; [
-    clr
-    hipblas
-    rocblas
-    llvm.lld
-    llvm.bintools
-  ];
+  rocmBuildInputs =
+    with rocmPackages;
+    [
+      clr
+      hipblas
+      rocblas
+      llvm.lld
+      llvm.bintools
+    ]
+    ++ optionals rocmUseWmma [ rocmPackages.rocwmma ];
 
   vulkanBuildInputs = [
     vulkan-headers
@@ -198,6 +202,10 @@ effectiveStdenv.mkDerivation (finalAttrs: {
       (cmakeFeature "AMDGPU_TARGETS" rocmGpuTargets)
       (cmakeBool "GGML_CUDA_FA_ALL_QUANTS" buildAllCudaFaQuants)
     ]
+    ++ optionals rocmUseWmma [
+      (cmakeBool "GGML_HIP_ROCWMMA_FATTN" rocmUseWmma)
+      (cmakeFeature "GGML_HIP_ROCWMMA_PATH" "${rocmPackages.rocwmma}")
+    ]
     ++ optionals useMetalKit [
       (lib.cmakeFeature "CMAKE_C_FLAGS" "-D__ARM_FEATURE_DOTPROD=1")
       (cmakeBool "GGML_METAL_EMBED_LIBRARY" (!precompileMetalShaders))

ggml/src/ggml-hip/CMakeLists.txt

@@ -39,7 +39,7 @@ endif()
 find_package(hip     REQUIRED)
 find_package(hipblas REQUIRED)
 find_package(rocblas REQUIRED)
-if (GGML_HIP_ROCWMMA_FATTN)
+if (GGML_HIP_ROCWMMA_FATTN AND NOT GGML_HIP_ROCWMMA_PATH)
     CHECK_INCLUDE_FILE_CXX("rocwmma/rocwmma.hpp" FOUND_ROCWMMA)
     if (NOT ${FOUND_ROCWMMA})
         message(FATAL_ERROR "rocwmma has not been found")
@@ -111,6 +111,9 @@ endif()
 
 if (GGML_HIP_ROCWMMA_FATTN)
     add_compile_definitions(GGML_HIP_ROCWMMA_FATTN)
+    if (GGML_HIP_ROCWMMA_PATH)
+        target_include_directories(ggml-hip PRIVATE ${GGML_HIP_ROCWMMA_PATH}/include)
+    endif()
 endif()
 
 if (NOT GGML_HIP_MMQ_MFMA)