From 04cfbd0435af59dc2321b5f727407358ae54240d Mon Sep 17 00:00:00 2001
From: Connor Baker <ConnorBaker01@gmail.com>
Date: Fri, 17 Oct 2025 07:19:18 +0000
Subject: [PATCH] cudaPackages.libcublasmp: use newly packaged libnvshmem

Signed-off-by: Connor Baker <ConnorBaker01@gmail.com>
---
 .../cuda-modules/packages/libcublasmp.nix            | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/pkgs/development/cuda-modules/packages/libcublasmp.nix b/pkgs/development/cuda-modules/packages/libcublasmp.nix
index 92130c135a06..8c0b7a5482f9 100644
--- a/pkgs/development/cuda-modules/packages/libcublasmp.nix
+++ b/pkgs/development/cuda-modules/packages/libcublasmp.nix
@@ -1,8 +1,11 @@
 {
   _cuda,
   buildRedist,
+  cuda_cudart,
+  lib,
   libcublas,
-  nvshmem ? null, # TODO(@connorbaker): package this
+  libnvshmem,
+  nccl,
 }:
 buildRedist {
   redistName = "cublasmp";
@@ -18,10 +21,15 @@ buildRedist {
   # TODO: Looks like the minimum supported capability is 7.0 as of the latest:
   # https://docs.nvidia.com/cuda/cublasmp/getting_started/index.html
   buildInputs = [
+    (lib.getOutput "stubs" cuda_cudart)
     libcublas
+    libnvshmem
+    nccl
   ];
 
-  platformAssertions = _cuda.lib._mkMissingPackagesAssertions { inherit nvshmem; };
+  extraAutoPatchelfLibs = [
+    "${lib.getOutput "stubs" cuda_cudart}/lib/stubs"
+  ];
 
   meta = {
     description = "High-performance, multi-process, GPU-accelerated library for distributed basic dense linear algebra";