Update docs

Author: github-actions[bot]

_sources/autoapi/tilelang/language/reduce/index.rst.txt

@@ -146,6 +146,30 @@ Module Contents
 
    Negative `dim` indices are normalized (Python-style). If `dst` is None, the operation is performed in-place into `src`. Raises ValueError when `dim` is out of bounds for `src.shape`. When `src.scope() == "local.fragment"`, this delegates to `cumsum_fragment`; otherwise it emits the `tl.cumsum` intrinsic.
 
+   .. rubric:: Examples
+
+   A 1D inclusive scan that writes the result into a separate shared-memory buffer:
+
+   >>> import tilelang.language as T
+   >>> @T.prim_func
+   ... def kernel(A: T.Tensor((128,), "float32"), B: T.Tensor((128,), "float32")):
+   ...     with T.Kernel(1, threads=128):
+   ...         A_shared = T.alloc_shared((128,), "float32")
+   ...         T.copy(A, A_shared)
+   ...         T.cumsum(src=A_shared, dst=A_shared, dim=0)
+   ...         T.copy(A_shared, B)
+
+   A 2D prefix sum along the last dimension with reverse accumulation:
+
+   >>> import tilelang.language as T
+   >>> @T.prim_func
+   ... def kernel2d(A: T.Tensor((64, 64), "float16"), B: T.Tensor((64, 64), "float16")):
+   ...     with T.Kernel(1, 1, threads=256):
+   ...         tile = T.alloc_shared((64, 64), "float16")
+   ...         T.copy(A, tile)
+   ...         T.cumsum(src=tile, dim=1, reverse=True)
+   ...         T.copy(tile, B)
+
    :returns: A handle to the emitted cumulative-sum operation.
    :rtype: tir.Call
 

autoapi/tilelang/language/reduce/index.html

@@ -695,6 +695,29 @@ <h2>Module Contents<a class="headerlink" href="#module-contents" title="Link to
 <span class="sig-prename descclassname"><span class="pre">tilelang.language.reduce.</span></span><span class="sig-name descname"><span class="pre">cumsum</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">src</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">dst</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">dim</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">reverse</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">False</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#tilelang.language.reduce.cumsum" title="Link to this definition">¶</a></dt>
 <dd><p>Compute the cumulative sum of <cite>src</cite> along <cite>dim</cite>, writing results to <cite>dst</cite>.</p>
 <p>Negative <cite>dim</cite> indices are normalized (Python-style). If <cite>dst</cite> is None, the operation is performed in-place into <cite>src</cite>. Raises ValueError when <cite>dim</cite> is out of bounds for <cite>src.shape</cite>. When <cite>src.scope() == “local.fragment”</cite>, this delegates to <cite>cumsum_fragment</cite>; otherwise it emits the <cite>tl.cumsum</cite> intrinsic.</p>
+<p class="rubric">Examples</p>
+<p>A 1D inclusive scan that writes the result into a separate shared-memory buffer:</p>
+<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">import</span><span class="w"> </span><span class="nn">tilelang.language</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">T</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="nd">@T</span><span class="o">.</span><span class="n">prim_func</span>
+<span class="gp">... </span><span class="k">def</span><span class="w"> </span><span class="nf">kernel</span><span class="p">(</span><span class="n">A</span><span class="p">:</span> <span class="n">T</span><span class="o">.</span><span class="n">Tensor</span><span class="p">((</span><span class="mi">128</span><span class="p">,),</span> <span class="s2">&quot;float32&quot;</span><span class="p">),</span> <span class="n">B</span><span class="p">:</span> <span class="n">T</span><span class="o">.</span><span class="n">Tensor</span><span class="p">((</span><span class="mi">128</span><span class="p">,),</span> <span class="s2">&quot;float32&quot;</span><span class="p">)):</span>
+<span class="gp">... </span>    <span class="k">with</span> <span class="n">T</span><span class="o">.</span><span class="n">Kernel</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">threads</span><span class="o">=</span><span class="mi">128</span><span class="p">):</span>
+<span class="gp">... </span>        <span class="n">A_shared</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">alloc_shared</span><span class="p">((</span><span class="mi">128</span><span class="p">,),</span> <span class="s2">&quot;float32&quot;</span><span class="p">)</span>
+<span class="gp">... </span>        <span class="n">T</span><span class="o">.</span><span class="n">copy</span><span class="p">(</span><span class="n">A</span><span class="p">,</span> <span class="n">A_shared</span><span class="p">)</span>
+<span class="gp">... </span>        <span class="n">T</span><span class="o">.</span><span class="n">cumsum</span><span class="p">(</span><span class="n">src</span><span class="o">=</span><span class="n">A_shared</span><span class="p">,</span> <span class="n">dst</span><span class="o">=</span><span class="n">A_shared</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
+<span class="gp">... </span>        <span class="n">T</span><span class="o">.</span><span class="n">copy</span><span class="p">(</span><span class="n">A_shared</span><span class="p">,</span> <span class="n">B</span><span class="p">)</span>
+</pre></div>
+</div>
+<p>A 2D prefix sum along the last dimension with reverse accumulation:</p>
+<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">import</span><span class="w"> </span><span class="nn">tilelang.language</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">T</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="nd">@T</span><span class="o">.</span><span class="n">prim_func</span>
+<span class="gp">... </span><span class="k">def</span><span class="w"> </span><span class="nf">kernel2d</span><span class="p">(</span><span class="n">A</span><span class="p">:</span> <span class="n">T</span><span class="o">.</span><span class="n">Tensor</span><span class="p">((</span><span class="mi">64</span><span class="p">,</span> <span class="mi">64</span><span class="p">),</span> <span class="s2">&quot;float16&quot;</span><span class="p">),</span> <span class="n">B</span><span class="p">:</span> <span class="n">T</span><span class="o">.</span><span class="n">Tensor</span><span class="p">((</span><span class="mi">64</span><span class="p">,</span> <span class="mi">64</span><span class="p">),</span> <span class="s2">&quot;float16&quot;</span><span class="p">)):</span>
+<span class="gp">... </span>    <span class="k">with</span> <span class="n">T</span><span class="o">.</span><span class="n">Kernel</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">threads</span><span class="o">=</span><span class="mi">256</span><span class="p">):</span>
+<span class="gp">... </span>        <span class="n">tile</span> <span class="o">=</span> <span class="n">T</span><span class="o">.</span><span class="n">alloc_shared</span><span class="p">((</span><span class="mi">64</span><span class="p">,</span> <span class="mi">64</span><span class="p">),</span> <span class="s2">&quot;float16&quot;</span><span class="p">)</span>
+<span class="gp">... </span>        <span class="n">T</span><span class="o">.</span><span class="n">copy</span><span class="p">(</span><span class="n">A</span><span class="p">,</span> <span class="n">tile</span><span class="p">)</span>
+<span class="gp">... </span>        <span class="n">T</span><span class="o">.</span><span class="n">cumsum</span><span class="p">(</span><span class="n">src</span><span class="o">=</span><span class="n">tile</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">reverse</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
+<span class="gp">... </span>        <span class="n">T</span><span class="o">.</span><span class="n">copy</span><span class="p">(</span><span class="n">tile</span><span class="p">,</span> <span class="n">B</span><span class="p">)</span>
+</pre></div>
+</div>
 <dl class="field-list simple">
 <dt class="field-odd">Returns<span class="colon">:</span></dt>
 <dd class="field-odd"><p>A handle to the emitted cumulative-sum operation.</p>