Update docs

Author: github-actions[bot]

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

@@ -18,13 +18,17 @@ Functions
    tilelang.language.customize.buffer_to_tile_region
    tilelang.language.customize.buffer_load_to_tile_region
    tilelang.language.customize.buffer_region_to_tile_region
+   tilelang.language.customize.atomic_max
+   tilelang.language.customize.atomic_min
    tilelang.language.customize.atomic_add
    tilelang.language.customize.atomic_addx2
    tilelang.language.customize.atomic_addx4
    tilelang.language.customize.dp4a
    tilelang.language.customize.clamp
    tilelang.language.customize.reshape
    tilelang.language.customize.view
+   tilelang.language.customize.atomic_load
+   tilelang.language.customize.atomic_store
 
 
 Module Contents
@@ -86,7 +90,33 @@ Module Contents
    :rtype: tir.Call
 
 
-.. py:function:: atomic_add(dst, value)
+.. py:function:: atomic_max(dst, value, memory_order = None)
+
+   Perform an atomic maximum operation.
+
+   :param dst: Destination buffer where the atomic maximum will be performed
+   :type dst: Buffer
+   :param value: Value to be atomically added
+   :type value: PrimExpr
+
+   :returns: Handle to the atomic maximum operation
+   :rtype: PrimExpr
+
+
+.. py:function:: atomic_min(dst, value, memory_order = None)
+
+   Perform an atomic minimum operation.
+
+   :param dst: Destination buffer where the atomic minimum will be performed
+   :type dst: Buffer
+   :param value: Value to be atomically added
+   :type value: PrimExpr
+
+   :returns: Handle to the atomic minimum operation
+   :rtype: PrimExpr
+
+
+.. py:function:: atomic_add(dst, value, memory_order = None)
 
    Perform an atomic addition operation.
 
@@ -179,3 +209,31 @@ Module Contents
    :rtype: Buffer
 
 
+.. py:function:: atomic_load(src, memory_order = 'seq_cst')
+
+   Loads a value from the input buffer with specified memory_order.
+
+   :param src: Input buffer to load from
+   :type src: Buffer
+   :param memory_order: Atomicity level for the load operation. Defaults to "seq_cst".
+   :type memory_order: str, optional
+
+   :returns: The loaded value from the buffer
+   :rtype: PrimExpr
+
+
+.. py:function:: atomic_store(dst, src, memory_order = 'seq_cst')
+
+   Stores a value to the input buffer with specified memory_order.
+
+   :param dst: Input buffer to store to
+   :type dst: Buffer
+   :param src: Value to store
+   :type src: PrimExpr
+   :param memory_order: Atomicity level for the load operation. Defaults to "seq_cst".
+   :type memory_order: str, optional
+
+   :returns: The handle of the store operation
+   :rtype: PrimExpr
+
+

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

@@ -22,6 +22,7 @@ Functions
    tilelang.language.reduce.reduce_absmax
    tilelang.language.reduce.cumsum_fragment
    tilelang.language.reduce.cumsum
+   tilelang.language.reduce.finalize_reducer
 
 
 Module Contents
@@ -156,3 +157,14 @@ Module Contents
    :rtype: tir.Call
 
 
+.. py:function:: finalize_reducer(reducer)
+
+   Finalize the reducer buffer.
+
+   :param reducer: The reducer buffer
+   :type reducer: tir.Buffer
+
+   :returns: Handle to the finalize reducer operation
+   :rtype: tir.Call
+
+

_sources/autoapi/tilelang/transform/index.rst.txt

@@ -63,6 +63,7 @@ Functions
    tilelang.transform.LowerOpaqueBlock
    tilelang.transform.LowerThreadAllreduce
    tilelang.transform.LowerDeviceKernelLaunch
+   tilelang.transform.LayoutReducer
 
 
 Package Contents
@@ -375,3 +376,9 @@ Package Contents
 
 
 
+.. py:function:: LayoutReducer()
+
+   LayoutReducer
+
+
+

autoapi/tilelang/language/customize/index.html

@@ -478,7 +478,13 @@ <h2>Functions<a class="headerlink" href="#functions" title="Link to this heading
 <tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.customize.buffer_region_to_tile_region" title="tilelang.language.customize.buffer_region_to_tile_region"><code class="xref py py-obj docutils literal notranslate"><span class="pre">buffer_region_to_tile_region</span></code></a>(buffer_region, ...)</p></td>
 <td><p>Convert a buffer region to a tile region descriptor.</p></td>
 </tr>
-<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.customize.atomic_add" title="tilelang.language.customize.atomic_add"><code class="xref py py-obj docutils literal notranslate"><span class="pre">atomic_add</span></code></a>(dst, value)</p></td>
+<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.customize.atomic_max" title="tilelang.language.customize.atomic_max"><code class="xref py py-obj docutils literal notranslate"><span class="pre">atomic_max</span></code></a>(dst, value[, memory_order])</p></td>
+<td><p>Perform an atomic maximum operation.</p></td>
+</tr>
+<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.customize.atomic_min" title="tilelang.language.customize.atomic_min"><code class="xref py py-obj docutils literal notranslate"><span class="pre">atomic_min</span></code></a>(dst, value[, memory_order])</p></td>
+<td><p>Perform an atomic minimum operation.</p></td>
+</tr>
+<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.customize.atomic_add" title="tilelang.language.customize.atomic_add"><code class="xref py py-obj docutils literal notranslate"><span class="pre">atomic_add</span></code></a>(dst, value[, memory_order])</p></td>
 <td><p>Perform an atomic addition operation.</p></td>
 </tr>
 <tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.customize.atomic_addx2" title="tilelang.language.customize.atomic_addx2"><code class="xref py py-obj docutils literal notranslate"><span class="pre">atomic_addx2</span></code></a>(dst, value)</p></td>
@@ -499,6 +505,12 @@ <h2>Functions<a class="headerlink" href="#functions" title="Link to this heading
 <tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.customize.view" title="tilelang.language.customize.view"><code class="xref py py-obj docutils literal notranslate"><span class="pre">view</span></code></a>(src[, shape, dtype])</p></td>
 <td><p>Views the input buffer with optionally modified shape and dtype.</p></td>
 </tr>
+<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.customize.atomic_load" title="tilelang.language.customize.atomic_load"><code class="xref py py-obj docutils literal notranslate"><span class="pre">atomic_load</span></code></a>(src[, memory_order])</p></td>
+<td><p>Loads a value from the input buffer with specified memory_order.</p></td>
+</tr>
+<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.customize.atomic_store" title="tilelang.language.customize.atomic_store"><code class="xref py py-obj docutils literal notranslate"><span class="pre">atomic_store</span></code></a>(dst, src[, memory_order])</p></td>
+<td><p>Stores a value to the input buffer with specified memory_order.</p></td>
+</tr>
 </tbody>
 </table>
 </div>
@@ -588,15 +600,58 @@ <h2>Module Contents<a class="headerlink" href="#module-contents" title="Link to
 </dl>
 </dd></dl>
 
+<dl class="py function">
+<dt class="sig sig-object py" id="tilelang.language.customize.atomic_max">
+<span class="sig-prename descclassname"><span class="pre">tilelang.language.customize.</span></span><span class="sig-name descname"><span class="pre">atomic_max</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">dst</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">value</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">memory_order</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#tilelang.language.customize.atomic_max" title="Link to this definition">¶</a></dt>
+<dd><p>Perform an atomic maximum operation.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters<span class="colon">:</span></dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>dst</strong> (<em>Buffer</em>) – Destination buffer where the atomic maximum will be performed</p></li>
+<li><p><strong>value</strong> (<em>PrimExpr</em>) – Value to be atomically added</p></li>
+<li><p><strong>memory_order</strong> (<em>str</em><em> | </em><em>None</em>)</p></li>
+</ul>
+</dd>
+<dt class="field-even">Returns<span class="colon">:</span></dt>
+<dd class="field-even"><p>Handle to the atomic maximum operation</p>
+</dd>
+<dt class="field-odd">Return type<span class="colon">:</span></dt>
+<dd class="field-odd"><p>PrimExpr</p>
+</dd>
+</dl>
+</dd></dl>
+
+<dl class="py function">
+<dt class="sig sig-object py" id="tilelang.language.customize.atomic_min">
+<span class="sig-prename descclassname"><span class="pre">tilelang.language.customize.</span></span><span class="sig-name descname"><span class="pre">atomic_min</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">dst</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">value</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">memory_order</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#tilelang.language.customize.atomic_min" title="Link to this definition">¶</a></dt>
+<dd><p>Perform an atomic minimum operation.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters<span class="colon">:</span></dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>dst</strong> (<em>Buffer</em>) – Destination buffer where the atomic minimum will be performed</p></li>
+<li><p><strong>value</strong> (<em>PrimExpr</em>) – Value to be atomically added</p></li>
+<li><p><strong>memory_order</strong> (<em>str</em><em> | </em><em>None</em>)</p></li>
+</ul>
+</dd>
+<dt class="field-even">Returns<span class="colon">:</span></dt>
+<dd class="field-even"><p>Handle to the atomic minimum operation</p>
+</dd>
+<dt class="field-odd">Return type<span class="colon">:</span></dt>
+<dd class="field-odd"><p>PrimExpr</p>
+</dd>
+</dl>
+</dd></dl>
+
 <dl class="py function">
 <dt class="sig sig-object py" id="tilelang.language.customize.atomic_add">
-<span class="sig-prename descclassname"><span class="pre">tilelang.language.customize.</span></span><span class="sig-name descname"><span class="pre">atomic_add</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">dst</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">value</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#tilelang.language.customize.atomic_add" title="Link to this definition">¶</a></dt>
+<span class="sig-prename descclassname"><span class="pre">tilelang.language.customize.</span></span><span class="sig-name descname"><span class="pre">atomic_add</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">dst</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">value</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">memory_order</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#tilelang.language.customize.atomic_add" title="Link to this definition">¶</a></dt>
 <dd><p>Perform an atomic addition operation.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters<span class="colon">:</span></dt>
 <dd class="field-odd"><ul class="simple">
 <li><p><strong>dst</strong> (<em>Buffer</em>) – Destination buffer where the atomic addition will be performed</p></li>
 <li><p><strong>value</strong> (<em>PrimExpr</em>) – Value to be atomically added</p></li>
+<li><p><strong>memory_order</strong> (<em>str</em><em> | </em><em>None</em>)</p></li>
 </ul>
 </dd>
 <dt class="field-even">Returns<span class="colon">:</span></dt>
@@ -731,6 +786,47 @@ <h2>Module Contents<a class="headerlink" href="#module-contents" title="Link to
 </dl>
 </dd></dl>
 
+<dl class="py function">
+<dt class="sig sig-object py" id="tilelang.language.customize.atomic_load">
+<span class="sig-prename descclassname"><span class="pre">tilelang.language.customize.</span></span><span class="sig-name descname"><span class="pre">atomic_load</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">memory_order</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">'seq_cst'</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#tilelang.language.customize.atomic_load" title="Link to this definition">¶</a></dt>
+<dd><p>Loads a value from the input buffer with specified memory_order.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters<span class="colon">:</span></dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>src</strong> (<em>Buffer</em>) – Input buffer to load from</p></li>
+<li><p><strong>memory_order</strong> (<em>str</em><em>, </em><em>optional</em>) – Atomicity level for the load operation. Defaults to “seq_cst”.</p></li>
+</ul>
+</dd>
+<dt class="field-even">Returns<span class="colon">:</span></dt>
+<dd class="field-even"><p>The loaded value from the buffer</p>
+</dd>
+<dt class="field-odd">Return type<span class="colon">:</span></dt>
+<dd class="field-odd"><p>PrimExpr</p>
+</dd>
+</dl>
+</dd></dl>
+
+<dl class="py function">
+<dt class="sig sig-object py" id="tilelang.language.customize.atomic_store">
+<span class="sig-prename descclassname"><span class="pre">tilelang.language.customize.</span></span><span class="sig-name descname"><span class="pre">atomic_store</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">dst</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">src</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">memory_order</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">'seq_cst'</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#tilelang.language.customize.atomic_store" title="Link to this definition">¶</a></dt>
+<dd><p>Stores a value to the input buffer with specified memory_order.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters<span class="colon">:</span></dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>dst</strong> (<em>Buffer</em>) – Input buffer to store to</p></li>
+<li><p><strong>src</strong> (<em>PrimExpr</em>) – Value to store</p></li>
+<li><p><strong>memory_order</strong> (<em>str</em><em>, </em><em>optional</em>) – Atomicity level for the load operation. Defaults to “seq_cst”.</p></li>
+</ul>
+</dd>
+<dt class="field-even">Returns<span class="colon">:</span></dt>
+<dd class="field-even"><p>The handle of the store operation</p>
+</dd>
+<dt class="field-odd">Return type<span class="colon">:</span></dt>
+<dd class="field-odd"><p>PrimExpr</p>
+</dd>
+</dl>
+</dd></dl>
+
 </section>
 </section>
 
@@ -796,13 +892,17 @@ <h2>Module Contents<a class="headerlink" href="#module-contents" title="Link to
 <li><a class="reference internal" href="#tilelang.language.customize.buffer_to_tile_region"><code class="docutils literal notranslate"><span class="pre">buffer_to_tile_region()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.language.customize.buffer_load_to_tile_region"><code class="docutils literal notranslate"><span class="pre">buffer_load_to_tile_region()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.language.customize.buffer_region_to_tile_region"><code class="docutils literal notranslate"><span class="pre">buffer_region_to_tile_region()</span></code></a></li>
+<li><a class="reference internal" href="#tilelang.language.customize.atomic_max"><code class="docutils literal notranslate"><span class="pre">atomic_max()</span></code></a></li>
+<li><a class="reference internal" href="#tilelang.language.customize.atomic_min"><code class="docutils literal notranslate"><span class="pre">atomic_min()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.language.customize.atomic_add"><code class="docutils literal notranslate"><span class="pre">atomic_add()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.language.customize.atomic_addx2"><code class="docutils literal notranslate"><span class="pre">atomic_addx2()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.language.customize.atomic_addx4"><code class="docutils literal notranslate"><span class="pre">atomic_addx4()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.language.customize.dp4a"><code class="docutils literal notranslate"><span class="pre">dp4a()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.language.customize.clamp"><code class="docutils literal notranslate"><span class="pre">clamp()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.language.customize.reshape"><code class="docutils literal notranslate"><span class="pre">reshape()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.language.customize.view"><code class="docutils literal notranslate"><span class="pre">view()</span></code></a></li>
+<li><a class="reference internal" href="#tilelang.language.customize.atomic_load"><code class="docutils literal notranslate"><span class="pre">atomic_load()</span></code></a></li>
+<li><a class="reference internal" href="#tilelang.language.customize.atomic_store"><code class="docutils literal notranslate"><span class="pre">atomic_store()</span></code></a></li>
 </ul>
 </li>
 </ul>