mirror of
https://github.com/ml-explore/mlx.git
synced 2025-09-04 06:44:40 +08:00
[CUDA] Initial implementation of Convolution with cuDNN (#2385)
* Link with cuDNN * Initial implementation * Remove backend apis * Fix recording cudnn conv * More unused backend apis * Fix C++ conv tests * include cudnn as python dep * Install libcudnn9-dev-cuda-12 in CI * cudnn only accepts contiguous inputs * Switch to backend apis * Plan needs to be kept alive * Turn off tf32 * Add cache * Test the native cuda graph api * Set cudnn stream before execution * Make LRUCache more like a normal container * Do error check for cublas handle * Zero-initilizing array * Use tf32 for conv * Skip TestConv.test_torch_conv_2D test --------- Co-authored-by: Awni Hannun <awni@apple.com>
This commit is contained in:
@@ -16,7 +16,7 @@ rm "${repaired_wheel}"
|
||||
mlx_so="mlx/lib/libmlx.so"
|
||||
rpath=$(patchelf --print-rpath "${mlx_so}")
|
||||
base="\$ORIGIN/../../nvidia"
|
||||
rpath=$rpath:${base}/cublas/lib:${base}/cuda_nvrtc/lib
|
||||
rpath=$rpath:${base}/cublas/lib:${base}/cuda_nvrtc/lib:${base}/cudnn/lib
|
||||
patchelf --force-rpath --set-rpath "$rpath" "$mlx_so"
|
||||
python ../python/scripts/repair_record.py ${mlx_so}
|
||||
|
||||
|
Reference in New Issue
Block a user