rocmPackages.{rocshmem,hipsparselt}: init at 7.2.3 (#538426)

pkgs/development/rocm-modules/default.nix

@@ -117,6 +117,8 @@ let
 
       rccl = self.callPackage ./rccl { };
 
+      rocshmem = self.callPackage ./rocshmem { };
+
       hipcub = self.callPackage ./hipcub { };
 
       hipsparse = self.callPackage ./hipsparse { };
@@ -155,6 +157,8 @@ let
 
       hipblaslt = self.callPackage ./hipblaslt { };
 
+      hipsparselt = self.callPackage ./hipsparselt { };
+
       # hipTensor - Only supports GFX9
 
       composable_kernel_base = self.callPackage ./composable_kernel/base.nix { };

pkgs/development/rocm-modules/hipsparselt/default.nix

@@ -0,0 +1,123 @@
+{
+  lib,
+  stdenv,
+  fetchFromGitHub,
+  rocmUpdateScript,
+  cmake,
+  ninja,
+  git,
+  pkg-config,
+  rocm-cmake,
+  rocm-smi,
+  clr,
+  hipsparse,
+  openmp,
+  llvm,
+  msgpack-cxx,
+  libxml2,
+  zlib,
+  zstd,
+  python3,
+  python3Packages,
+  # hipsparselt only supports gfx942/gfx950
+  gpuTargets ? (clr.localGpuTargets or clr.gpuTargets),
+}:
+
+let
+  supportedTargets = (
+    lib.lists.intersectLists gpuTargets [
+      "gfx942"
+      "gfx950"
+    ]
+  );
+  supportsTargetArches = supportedTargets != [ ];
+  py = python3.withPackages (ps: [
+    ps.pyyaml
+    ps.setuptools
+    ps.packaging
+    ps.nanobind
+    ps.msgpack
+    ps.joblib
+  ]);
+  compiler = "amdclang++";
+in
+stdenv.mkDerivation (finalAttrs: {
+  pname = "hipsparselt${clr.gpuArchSuffix}";
+  version = "7.2.3";
+
+  src = fetchFromGitHub {
+    owner = "ROCm";
+    repo = "rocm-libraries";
+    rev = "rocm-${finalAttrs.version}";
+    hash = "sha256-Al/0yZPn/wqtBFQU1FcGoSquDMObk6HQ3D5sSh5biHg=";
+    sparseCheckout = [
+      "projects/hipsparselt"
+      "projects/hipblaslt/tensilelite"
+      "projects/hipblaslt/cmake"
+      "projects/hipblaslt/device-library"
+      "shared/origami"
+    ];
+  };
+  sourceRoot = "${finalAttrs.src.name}/projects/hipsparselt";
+
+  env.CXX = compiler;
+  env.ROCM_PATH = "${clr}";
+  env.TENSILE_ROCM_ASSEMBLER_PATH = lib.getExe' clr "amdclang++";
+  env.TENSILE_GEN_ASSEMBLY_TOOLCHAIN = lib.getExe' clr "amdclang++";
+  requiredSystemFeatures = [ "big-parallel" ];
+
+  __structuredAttrs = true;
+  strictDeps = true;
+
+  nativeBuildInputs = [
+    cmake
+    ninja
+    git
+    pkg-config
+    rocm-cmake
+    rocm-smi
+    py
+    clr
+  ];
+
+  buildInputs = [
+    llvm.llvm
+    clr
+    rocm-cmake
+    hipsparse
+    openmp
+    msgpack-cxx
+    libxml2
+    python3Packages.msgpack
+    zlib
+    zstd
+  ];
+
+  cmakeFlags = [
+    (lib.cmakeFeature "GPU_TARGETS" (lib.concatStringsSep ";" supportedTargets))
+    (lib.cmakeBool "HIPSPARSELT_ENABLE_DEVICE" supportsTargetArches)
+    (lib.cmakeBool "HIPSPARSELT_ENABLE_CLIENT" false)
+    (lib.cmakeBool "HIPSPARSELT_BUILD_TESTING" false)
+    (lib.cmakeBool "HIPSPARSELT_ENABLE_MARKER" false)
+    (lib.cmakeBool "HIPBLASLT_ENABLE_ROCROLLER" false)
+    (lib.cmakeBool "HIPBLASLT_ENABLE_MARKER" false)
+    (lib.cmakeFeature "CMAKE_C_COMPILER" "amdclang")
+    (lib.cmakeFeature "CMAKE_HIP_COMPILER" compiler)
+    (lib.cmakeFeature "CMAKE_CXX_COMPILER" compiler)
+    (lib.cmakeFeature "FETCHCONTENT_SOURCE_DIR_NANOBIND" "${python3Packages.nanobind.src}")
+    (lib.cmakeFeature "CMAKE_INSTALL_BINDIR" "bin")
+    (lib.cmakeFeature "CMAKE_INSTALL_LIBDIR" "lib")
+    (lib.cmakeFeature "CMAKE_INSTALL_INCLUDEDIR" "include")
+  ];
+
+  passthru.supportsTargetArches = supportsTargetArches;
+  passthru.updateScript = rocmUpdateScript { inherit finalAttrs; };
+
+  meta = {
+    description = "ROCm hipSPARSELt - a SPARSE marshalling library";
+    homepage = "https://github.com/ROCm/rocm-libraries/tree/develop/projects/hipsparselt";
+    license = with lib.licenses; [ mit ];
+    teams = [ lib.teams.rocm ];
+    platforms = lib.platforms.linux;
+  };
+})

pkgs/development/rocm-modules/rocshmem/default.nix

@@ -0,0 +1,70 @@
+{
+  lib,
+  stdenv,
+  fetchFromGitHub,
+  rocmUpdateScript,
+  cmake,
+  rocm-cmake,
+  rocm-core,
+  rocm-runtime,
+  clr,
+  mpi,
+  useMpi ? true,
+  useReverseOffload ? useMpi,
+  buildExamples ? false,
+  gpuTargets ? (clr.localGpuTargets or [ ]),
+}:
+
+stdenv.mkDerivation (finalAttrs: {
+  pname = "rocshmem${clr.gpuArchSuffix}";
+  version = "7.2.3";
+
+  src = fetchFromGitHub {
+    owner = "ROCm";
+    repo = "rocSHMEM";
+    rev = "rocm-${finalAttrs.version}";
+    hash = "sha256-7TiUABwCVPvosoLdR8zN+QGiX+OD7iD+HKMLnGMmRYA=";
+  };
+
+  nativeBuildInputs = [
+    cmake
+    rocm-cmake
+    clr
+  ];
+
+  buildInputs = [
+    rocm-runtime
+  ]
+  ++ lib.optionals useMpi [
+    mpi
+  ];
+
+  cmakeFlags = [
+    (lib.cmakeFeature "CMAKE_BUILD_TYPE" "Release")
+    (lib.cmakeBool "CMAKE_POSITION_INDEPENDENT_CODE" true)
+    (lib.cmakeFeature "CMAKE_CXX_COMPILER" "hipcc")
+    (lib.cmakeFeature "ROCM_PATH" "${clr}")
+    (lib.cmakeFeature "ROCM_VERSION" rocm-core.ROCM_LIBPATCH_VERSION)
+    (lib.cmakeBool "USE_IPC" true)
+    (lib.cmakeBool "USE_RO" useReverseOffload)
+    (lib.cmakeBool "USE_SINGLE_NODE" (!useMpi))
+    (lib.cmakeBool "USE_EXTERNAL_MPI" useMpi)
+    (lib.cmakeBool "BUILD_EXAMPLES" buildExamples)
+    (lib.cmakeFeature "CMAKE_INSTALL_BINDIR" "bin")
+    (lib.cmakeFeature "CMAKE_INSTALL_LIBDIR" "lib")
+    (lib.cmakeFeature "CMAKE_INSTALL_INCLUDEDIR" "include")
+  ]
+  ++ lib.optionals (gpuTargets != [ ]) [
+    (lib.cmakeFeature "GPU_TARGETS" (lib.concatStringsSep ";" gpuTargets))
+  ];
+
+  passthru.updateScript = rocmUpdateScript { inherit finalAttrs; };
+
+  meta = {
+    description = "The ROCm OpenSHMEM (rocSHMEM) runtime";
+    homepage = "https://github.com/ROCm/rocSHMEM";
+    license = with lib.licenses; [ mit ];
+    teams = [ lib.teams.rocm ];
+    platforms = lib.platforms.linux;
+  };
+})