[QUESTION] Warp on Jetson Orin NX - Install succeeds, but wp.init() fails

[sam-t-pratt]

I have been learning Warp over the past few weeks on an x86-64 system without issue - it's a great tool!

I am now trying to transition the code to a Jetson Orin NX to use in our systems. Warp installs without issue using warp_lang-1.2.2+cu11-py3-none-manylinux2014_aarch64.whl, but calling wp.init() fails with a RuntimeError:

RuntimeError: Failed to load the shared library '~/env/lib/python3.12/site-packages/warp/bin/warp.so'

I can see the library in site-packages, so that's not the issue. My installed CUDA version (from nvcc --version) is 11.4.315. Are there other specific system dependencies I should check?

Before I get set up to build Warp from source, I figured I'd check to see if: 1) I'm doing something dumb, and 2) This is possible (and thus worth building)

Thanks for your help!

[shi-eric]

Hey @sam-t-pratt, we've previously tested the ability to run Warp on Jetson AGX Orin and I think the Orin NX should be similar enough that we expect it to work. Can you start by pasting the output of nvidia-smi? Additionally there should be a CUDA Version in the header. I think this needs to be at least 11.5 per the installation requirements.

Here's what it looks like on my Jetson using the warp-lang-1.2.2+cu12 wheel from GitHub:

nvidia-smi
Mon Jul 15 16:57:09 2024       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 540.3.0                Driver Version: N/A          CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  Orin (nvgpu)                  N/A  | N/A              N/A |                  N/A |
| N/A   N/A  N/A               N/A /  N/A | Not Supported        |     N/A          N/A |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+

python3 -m warp.examples.core.example_graph_capture
Warp 1.2.2 initialized:
   CUDA Toolkit 12.5, Driver 12.2
   Devices:
     "cpu"      : "aarch64"
     "cuda:0"   : "Orin" (61 GiB, sm_87, mempool enabled)
   Kernel cache:
     /home/eshi/.cache/warp/1.2.2
Module warp.utils 48da1c2 load on device 'cuda:0' took 1.02 ms
Module __main__ f2f53ef load on device 'cuda:0' took 0.40 ms

[sam-t-pratt]

Hey @shi-eric, sorry for the delay in responding; the issue ended up being that my Orin was on the incorrect CUDA version as you mentioned, but I managed to nuke my board configuration while updating CUDA. Once I reflashed my board (amateur mistake..) I was able to get Warp running successfully. RTFM strikes again - thank you for the help.

[shi-eric]

Thanks for the follow up! Glad to hear you were able to get things working.