Skip to content

Commit

Permalink
Update documentations
Browse files Browse the repository at this point in the history
  • Loading branch information
@actions-user
actions-user committed Sep 7, 2024
1 parent 51bef54 commit 52389c4
Show file tree
Hide file tree
Showing 47 changed files with 1,903 additions and 160 deletions.
457 changes: 457 additions & 0 deletions _modules/hippynn/custom_kernels/env_triton.html
View file

Large diffs are not rendered by default.

8 changes: 6 additions & 2 deletions _modules/hippynn/custom_kernels/tensor_wrapper.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -124,8 +124,8 @@ <h1>Source code for hippynn.custom_kernels.tensor_wrapper</h1><div class="highli
<a class="viewcode-back" href="../../../api_documentation/hippynn.custom_kernels.tensor_wrapper.html#hippynn.custom_kernels.tensor_wrapper.NumbaCompatibleTensorFunction.__init__">[docs]</a>
<span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="k">if</span> <span class="n">numba</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">is_available</span><span class="p">():</span>
<span class="bp">self</span><span class="o">.</span><span class="n">kernel64</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">make_kernel</span><span class="p">(</span><span class="n">numba</span><span class="o">.</span><span class="n">float64</span><span class="p">)</span>
<span class="bp">self</span><span class="o">.</span><span class="n">kernel32</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">make_kernel</span><span class="p">(</span><span class="n">numba</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>
<span class="bp">self</span><span class="o">.</span><span class="n">kernel64</span> <span class="o">=</span> <span class="kc">None</span>
<span class="bp">self</span><span class="o">.</span><span class="n">kernel32</span> <span class="o">=</span> <span class="kc">None</span>
<span class="k">else</span><span class="p">:</span>
<span class="bp">self</span><span class="o">.</span><span class="n">kernel64</span> <span class="o">=</span> <span class="n">_numba_gpu_not_found</span>
<span class="bp">self</span><span class="o">.</span><span class="n">kernel32</span> <span class="o">=</span> <span class="n">_numba_gpu_not_found</span></div>
Expand All @@ -146,8 +146,12 @@ <h1>Source code for hippynn.custom_kernels.tensor_wrapper</h1><div class="highli
<span class="k">with</span> <span class="n">numba</span><span class="o">.</span><span class="n">cuda</span><span class="o">.</span><span class="n">gpus</span><span class="p">[</span><span class="n">dev</span><span class="o">.</span><span class="n">index</span><span class="p">]:</span>
<span class="n">numba_args</span> <span class="o">=</span> <span class="n">batch_convert_torch_to_numba</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">)</span>
<span class="k">if</span> <span class="n">dtype</span> <span class="o">==</span> <span class="n">torch</span><span class="o">.</span><span class="n">float64</span><span class="p">:</span>
<span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">kernel64</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
<span class="bp">self</span><span class="o">.</span><span class="n">kernel64</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">make_kernel</span><span class="p">(</span><span class="n">numba</span><span class="o">.</span><span class="n">float64</span><span class="p">)</span>
<span class="bp">self</span><span class="o">.</span><span class="n">kernel64</span><span class="p">[</span><span class="n">launch_bounds</span><span class="p">](</span><span class="o">*</span><span class="n">numba_args</span><span class="p">)</span>
<span class="k">elif</span> <span class="n">dtype</span> <span class="o">==</span> <span class="n">torch</span><span class="o">.</span><span class="n">float32</span><span class="p">:</span>
<span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">kernel32</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
<span class="bp">self</span><span class="o">.</span><span class="n">kernel32</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">make_kernel</span><span class="p">(</span><span class="n">numba</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>
<span class="bp">self</span><span class="o">.</span><span class="n">kernel32</span><span class="p">[</span><span class="n">launch_bounds</span><span class="p">](</span><span class="o">*</span><span class="n">numba_args</span><span class="p">)</span>
<span class="k">else</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Bad dtype: </span><span class="si">{}</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">dtype</span><span class="p">))</span>
Expand Down
7 changes: 7 additions & 0 deletions _modules/hippynn/custom_kernels/test_env_numba.html
View file
Original file line number Diff line number Diff line change
Expand Up @@ -204,6 +204,7 @@ <h1>Source code for hippynn.custom_kernels.test_env_numba</h1><div class="highli
<span class="n">TEST_LARGE_PARAMS</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(</span><span class="n">n_molecules</span><span class="o">=</span><span class="mi">1000</span><span class="p">,</span> <span class="n">n_atoms</span><span class="o">=</span><span class="mi">30</span><span class="p">,</span> <span class="n">atom_prob</span><span class="o">=</span><span class="mf">0.7</span><span class="p">,</span> <span class="n">n_features</span><span class="o">=</span><span class="mi">80</span><span class="p">,</span> <span class="n">n_nu</span><span class="o">=</span><span class="mi">20</span><span class="p">)</span>
<span class="n">TEST_MEGA_PARAMS</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(</span><span class="n">n_molecules</span><span class="o">=</span><span class="mi">500</span><span class="p">,</span> <span class="n">n_atoms</span><span class="o">=</span><span class="mi">30</span><span class="p">,</span> <span class="n">atom_prob</span><span class="o">=</span><span class="mf">0.7</span><span class="p">,</span> <span class="n">n_features</span><span class="o">=</span><span class="mi">128</span><span class="p">,</span> <span class="n">n_nu</span><span class="o">=</span><span class="mi">100</span><span class="p">)</span>
<span class="n">TEST_ULTRA_PARAMS</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(</span><span class="n">n_molecules</span><span class="o">=</span><span class="mi">500</span><span class="p">,</span> <span class="n">n_atoms</span><span class="o">=</span><span class="mi">30</span><span class="p">,</span> <span class="n">atom_prob</span><span class="o">=</span><span class="mf">0.7</span><span class="p">,</span> <span class="n">n_features</span><span class="o">=</span><span class="mi">128</span><span class="p">,</span> <span class="n">n_nu</span><span class="o">=</span><span class="mi">320</span><span class="p">)</span>
<span class="n">TEST_GIGA_PARAMS</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(</span><span class="n">n_molecules</span><span class="o">=</span><span class="mi">32</span><span class="p">,</span> <span class="n">n_atoms</span><span class="o">=</span><span class="mi">30</span><span class="p">,</span> <span class="n">atom_prob</span><span class="o">=</span><span class="mf">0.7</span><span class="p">,</span> <span class="n">n_features</span><span class="o">=</span><span class="mi">512</span><span class="p">,</span> <span class="n">n_nu</span><span class="o">=</span><span class="mi">320</span><span class="p">)</span>

<span class="c1"># reference implementation</span>

Expand Down Expand Up @@ -566,6 +567,12 @@ <h1>Source code for hippynn.custom_kernels.test_env_numba</h1><div class="highli

<span class="k">if</span> <span class="n">use_verylarge_gpu</span><span class="p">:</span>
<span class="k">if</span> <span class="n">use_ultra</span><span class="p">:</span>

<span class="nb">print</span><span class="p">(</span><span class="s2">&quot;-&quot;</span> <span class="o">*</span> <span class="mi">80</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="s2">&quot;Giga systems:&quot;</span><span class="p">,</span> <span class="n">TEST_GIGA_PARAMS</span><span class="p">)</span>
<span class="n">tester</span><span class="o">.</span><span class="n">check_speed</span><span class="p">(</span>
<span class="n">n_repetitions</span><span class="o">=</span><span class="mi">20</span><span class="p">,</span> <span class="n">data_size</span><span class="o">=</span><span class="n">TEST_GIGA_PARAMS</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">device</span><span class="p">(</span><span class="s2">&quot;cuda&quot;</span><span class="p">),</span> <span class="n">compare_against</span><span class="o">=</span><span class="n">compare_against</span>
<span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="s2">&quot;-&quot;</span> <span class="o">*</span> <span class="mi">80</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="s2">&quot;Ultra systems:&quot;</span><span class="p">,</span> <span class="n">TEST_ULTRA_PARAMS</span><span class="p">)</span>
<span class="n">tester</span><span class="o">.</span><span class="n">check_speed</span><span class="p">(</span>
Expand Down
Loading

0 comments on commit 52389c4

Please sign in to comment.