|
399 | 399 | <li class="toctree-l1"><a class="reference internal" href="../notes/gradcheck.html">Gradcheck mechanics</a></li> |
400 | 400 | <li class="toctree-l1"><a class="reference internal" href="../notes/hip.html">HIP (ROCm) semantics</a></li> |
401 | 401 | <li class="toctree-l1"><a class="reference internal" href="../notes/large_scale_deployments.html">Features for large-scale deployments</a></li> |
| 402 | +<li class="toctree-l1"><a class="reference internal" href="../notes/libtorch_stable_abi.html">LibTorch Stable ABI</a></li> |
402 | 403 | <li class="toctree-l1"><a class="reference internal" href="../notes/modules.html">Modules</a></li> |
403 | 404 | <li class="toctree-l1"><a class="reference internal" href="../notes/mps.html">MPS backend</a></li> |
404 | 405 | <li class="toctree-l1"><a class="reference internal" href="../notes/multiprocessing.html">Multiprocessing best practices</a></li> |
@@ -860,28 +861,26 @@ <h1>Source code for torch</h1><div class="highlight"><pre> |
860 | 861 | <span class="k">del</span> <span class="n">_load_dll_libraries</span> |
861 | 862 |
|
862 | 863 |
|
| 864 | +<span class="k">def</span> <span class="nf">_get_cuda_dep_paths</span><span class="p">(</span><span class="n">path</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">lib_folder</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">lib_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-></span> <span class="nb">list</span><span class="p">[</span><span class="nb">str</span><span class="p">]:</span> |
| 865 | + <span class="c1"># Libraries can either be in path/nvidia/lib_folder/lib or path/lib_folder/lib</span> |
| 866 | + <span class="n">nvidia_lib_paths</span> <span class="o">=</span> <span class="n">glob</span><span class="o">.</span><span class="n">glob</span><span class="p">(</span> |
| 867 | + <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">path</span><span class="p">,</span> <span class="s2">"nvidia"</span><span class="p">,</span> <span class="n">lib_folder</span><span class="p">,</span> <span class="s2">"lib"</span><span class="p">,</span> <span class="n">lib_name</span><span class="p">)</span> |
| 868 | + <span class="p">)</span> |
| 869 | + <span class="n">lib_paths</span> <span class="o">=</span> <span class="n">glob</span><span class="o">.</span><span class="n">glob</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">path</span><span class="p">,</span> <span class="n">lib_folder</span><span class="p">,</span> <span class="s2">"lib"</span><span class="p">,</span> <span class="n">lib_name</span><span class="p">))</span> |
| 870 | + |
| 871 | + <span class="k">return</span> <span class="n">nvidia_lib_paths</span> <span class="o">+</span> <span class="n">lib_paths</span> |
| 872 | + |
| 873 | + |
863 | 874 | <span class="k">def</span> <span class="nf">_preload_cuda_deps</span><span class="p">(</span><span class="n">lib_folder</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">lib_name</span><span class="p">:</span> <span class="nb">str</span><span class="p">)</span> <span class="o">-></span> <span class="kc">None</span><span class="p">:</span> |
864 | 875 | <span class="w"> </span><span class="sd">"""Preloads cuda deps if they could not be found otherwise."""</span> |
865 | 876 | <span class="c1"># Should only be called on Linux if default path resolution have failed</span> |
866 | 877 | <span class="k">assert</span> <span class="n">platform</span><span class="o">.</span><span class="n">system</span><span class="p">()</span> <span class="o">==</span> <span class="s2">"Linux"</span><span class="p">,</span> <span class="s2">"Should only be called on Linux"</span> |
867 | 878 |
|
868 | 879 | <span class="n">lib_path</span> <span class="o">=</span> <span class="kc">None</span> |
869 | 880 | <span class="k">for</span> <span class="n">path</span> <span class="ow">in</span> <span class="n">sys</span><span class="o">.</span><span class="n">path</span><span class="p">:</span> |
870 | | - <span class="n">nvidia_path</span> <span class="o">=</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">path</span><span class="p">,</span> <span class="s2">"nvidia"</span><span class="p">)</span> |
871 | | - <span class="k">if</span> <span class="ow">not</span> <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">exists</span><span class="p">(</span><span class="n">nvidia_path</span><span class="p">):</span> |
872 | | - <span class="k">continue</span> |
873 | | - <span class="n">candidate_lib_paths</span> <span class="o">=</span> <span class="n">glob</span><span class="o">.</span><span class="n">glob</span><span class="p">(</span> |
874 | | - <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">nvidia_path</span><span class="p">,</span> <span class="n">lib_folder</span><span class="p">,</span> <span class="s2">"lib"</span><span class="p">,</span> <span class="n">lib_name</span><span class="p">)</span> |
875 | | - <span class="p">)</span> |
876 | | - <span class="c1"># if path/nvidia/lib_folder/ is not found look in path/lib_folder/</span> |
877 | | - <span class="k">if</span> <span class="ow">not</span> <span class="n">candidate_lib_paths</span><span class="p">:</span> |
878 | | - <span class="n">candidate_lib_paths</span> <span class="o">=</span> <span class="n">glob</span><span class="o">.</span><span class="n">glob</span><span class="p">(</span> |
879 | | - <span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">path</span><span class="p">,</span> <span class="n">lib_folder</span><span class="p">,</span> <span class="s2">"lib"</span><span class="p">,</span> <span class="n">lib_name</span><span class="p">)</span> |
880 | | - <span class="p">)</span> |
881 | | - |
882 | | - <span class="k">if</span> <span class="n">candidate_lib_paths</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">lib_path</span><span class="p">:</span> |
| 881 | + <span class="n">candidate_lib_paths</span> <span class="o">=</span> <span class="n">_get_cuda_dep_paths</span><span class="p">(</span><span class="n">path</span><span class="p">,</span> <span class="n">lib_folder</span><span class="p">,</span> <span class="n">lib_name</span><span class="p">)</span> |
| 882 | + <span class="k">if</span> <span class="n">candidate_lib_paths</span><span class="p">:</span> |
883 | 883 | <span class="n">lib_path</span> <span class="o">=</span> <span class="n">candidate_lib_paths</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> |
884 | | - <span class="k">if</span> <span class="n">lib_path</span><span class="p">:</span> |
885 | 884 | <span class="k">break</span> |
886 | 885 | <span class="k">if</span> <span class="ow">not</span> <span class="n">lib_path</span><span class="p">:</span> |
887 | 886 | <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">lib_name</span><span class="si">}</span><span class="s2"> not found in the system path </span><span class="si">{</span><span class="n">sys</span><span class="o">.</span><span class="n">path</span><span class="si">}</span><span class="s2">"</span><span class="p">)</span> |
|
0 commit comments