Update docs

Author: github-actions[bot]

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

@@ -244,6 +244,9 @@ Module Contents
    .. py:method:: disable_cache()
 
 
+   .. py:method:: is_print_on_compilation_enabled()
+
+
 .. py:data:: env
 
 .. py:data:: CUDA_HOME

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

@@ -24,7 +24,10 @@ Functions
    tilelang.language.builtin.set_max_nreg
    tilelang.language.builtin.inc_max_nreg
    tilelang.language.builtin.dec_max_nreg
+   tilelang.language.builtin.annotate_producer_reg_dealloc
+   tilelang.language.builtin.annotate_consumer_reg_alloc
    tilelang.language.builtin.no_set_max_nreg
+   tilelang.language.builtin.disable_warp_group_reg_alloc
    tilelang.language.builtin.mbarrier_wait_parity
    tilelang.language.builtin.mbarrier_arrive
    tilelang.language.builtin.mbarrier_expect_tx
@@ -149,12 +152,30 @@ Module Contents
 
 
 
+.. py:function:: annotate_producer_reg_dealloc(reg_count = 24)
+
+   Annotate the producer reg dealloc.
+
+
+
+.. py:function:: annotate_consumer_reg_alloc(reg_count = 240)
+
+   Annotate the consumer reg alloc.
+
+
+
 .. py:function:: no_set_max_nreg()
 
    Disable the maximum register limit setting.
 
 
 
+.. py:function:: disable_warp_group_reg_alloc()
+
+   Disable the warp group reg alloc.
+
+
+
 .. py:function:: mbarrier_wait_parity(mbarrier, parity)
 
    Wait for memory barrier parity condition.

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

@@ -40,6 +40,7 @@ Functions
    tilelang.transform.MergeIfStmt
    tilelang.transform.MultiVersionBuffer
    tilelang.transform.WarpSpecialized
+   tilelang.transform.AnnotateWarpGroupRegAlloc
    tilelang.transform.InjectTmaBarrier
    tilelang.transform.InjectFenceProxy
    tilelang.transform.LegalizeVectorizedLoop
@@ -198,6 +199,18 @@ Package Contents
    :rtype: tvm.transform.Pass
 
 
+.. py:function:: AnnotateWarpGroupRegAlloc()
+
+   Inject set_max_nreg calls into warp-specialized functions.
+
+   This pass analyzes the function to collect register hints from set_max_nreg
+   and no_set_max_nreg calls, then injects appropriate set_max_nreg calls into
+   producer and consumer branches of warp-specialized code.
+
+   :returns: **fpass** -- The result pass
+   :rtype: tvm.transform.Pass
+
+
 .. py:function:: InjectTmaBarrier()
 
    InjectTmaBarrier

autoapi/tilelang/env/index.html

@@ -831,6 +831,16 @@ <h3>Benefits<a class="headerlink" href="#benefits" title="Link to this heading">
 </dl>
 </dd></dl>
 
+<dl class="py method">
+<dt class="sig sig-object py" id="tilelang.env.Environment.is_print_on_compilation_enabled">
+<span class="sig-name descname"><span class="pre">is_print_on_compilation_enabled</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="headerlink" href="#tilelang.env.Environment.is_print_on_compilation_enabled" title="Link to this definition">¶</a></dt>
+<dd><dl class="field-list simple">
+<dt class="field-odd">Return type<span class="colon">:</span></dt>
+<dd class="field-odd"><p>bool</p>
+</dd>
+</dl>
+</dd></dl>
+
 </dd></dl>
 
 <dl class="py data">
@@ -988,6 +998,7 @@ <h3>Benefits<a class="headerlink" href="#benefits" title="Link to this heading">
 <li><a class="reference internal" href="#tilelang.env.Environment.is_cache_enabled"><code class="docutils literal notranslate"><span class="pre">Environment.is_cache_enabled()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.env.Environment.enable_cache"><code class="docutils literal notranslate"><span class="pre">Environment.enable_cache()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.env.Environment.disable_cache"><code class="docutils literal notranslate"><span class="pre">Environment.disable_cache()</span></code></a></li>
+<li><a class="reference internal" href="#tilelang.env.Environment.is_print_on_compilation_enabled"><code class="docutils literal notranslate"><span class="pre">Environment.is_print_on_compilation_enabled()</span></code></a></li>
 </ul>
 </li>
 <li><a class="reference internal" href="#tilelang.env.env"><code class="docutils literal notranslate"><span class="pre">env</span></code></a></li>

autoapi/tilelang/language/builtin/index.html

@@ -496,46 +496,55 @@ <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.builtin.dec_max_nreg" title="tilelang.language.builtin.dec_max_nreg"><code class="xref py py-obj docutils literal notranslate"><span class="pre">dec_max_nreg</span></code></a>(reg_count)</p></td>
 <td><p>Decrement the maximum number of registers to use.</p></td>
 </tr>
+<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.builtin.annotate_producer_reg_dealloc" title="tilelang.language.builtin.annotate_producer_reg_dealloc"><code class="xref py py-obj docutils literal notranslate"><span class="pre">annotate_producer_reg_dealloc</span></code></a>([reg_count])</p></td>
+<td><p>Annotate the producer reg dealloc.</p></td>
+</tr>
+<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.builtin.annotate_consumer_reg_alloc" title="tilelang.language.builtin.annotate_consumer_reg_alloc"><code class="xref py py-obj docutils literal notranslate"><span class="pre">annotate_consumer_reg_alloc</span></code></a>([reg_count])</p></td>
+<td><p>Annotate the consumer reg alloc.</p></td>
+</tr>
 <tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.builtin.no_set_max_nreg" title="tilelang.language.builtin.no_set_max_nreg"><code class="xref py py-obj docutils literal notranslate"><span class="pre">no_set_max_nreg</span></code></a>()</p></td>
 <td><p>Disable the maximum register limit setting.</p></td>
 </tr>
-<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.builtin.mbarrier_wait_parity" title="tilelang.language.builtin.mbarrier_wait_parity"><code class="xref py py-obj docutils literal notranslate"><span class="pre">mbarrier_wait_parity</span></code></a>(mbarrier, parity)</p></td>
+<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.builtin.disable_warp_group_reg_alloc" title="tilelang.language.builtin.disable_warp_group_reg_alloc"><code class="xref py py-obj docutils literal notranslate"><span class="pre">disable_warp_group_reg_alloc</span></code></a>()</p></td>
+<td><p>Disable the warp group reg alloc.</p></td>
+</tr>
+<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.builtin.mbarrier_wait_parity" title="tilelang.language.builtin.mbarrier_wait_parity"><code class="xref py py-obj docutils literal notranslate"><span class="pre">mbarrier_wait_parity</span></code></a>(mbarrier, parity)</p></td>
 <td><p>Wait for memory barrier parity condition.</p></td>
 </tr>
-<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.builtin.mbarrier_arrive" title="tilelang.language.builtin.mbarrier_arrive"><code class="xref py py-obj docutils literal notranslate"><span class="pre">mbarrier_arrive</span></code></a>(mbarrier)</p></td>
+<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.builtin.mbarrier_arrive" title="tilelang.language.builtin.mbarrier_arrive"><code class="xref py py-obj docutils literal notranslate"><span class="pre">mbarrier_arrive</span></code></a>(mbarrier)</p></td>
 <td><p>Arrive at memory barrier.</p></td>
 </tr>
-<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.builtin.mbarrier_expect_tx" title="tilelang.language.builtin.mbarrier_expect_tx"><code class="xref py py-obj docutils literal notranslate"><span class="pre">mbarrier_expect_tx</span></code></a>(*args)</p></td>
+<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.builtin.mbarrier_expect_tx" title="tilelang.language.builtin.mbarrier_expect_tx"><code class="xref py py-obj docutils literal notranslate"><span class="pre">mbarrier_expect_tx</span></code></a>(*args)</p></td>
 <td><p>Set expected transaction count for memory barrier.</p></td>
 </tr>
-<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.builtin.wait_wgmma" title="tilelang.language.builtin.wait_wgmma"><code class="xref py py-obj docutils literal notranslate"><span class="pre">wait_wgmma</span></code></a>(id)</p></td>
+<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.builtin.wait_wgmma" title="tilelang.language.builtin.wait_wgmma"><code class="xref py py-obj docutils literal notranslate"><span class="pre">wait_wgmma</span></code></a>(id)</p></td>
 <td><p>Wait for WGMMA (Warp Group Matrix Multiply-Accumulate) operations to complete.</p></td>
 </tr>
-<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.builtin.barrier_wait" title="tilelang.language.builtin.barrier_wait"><code class="xref py py-obj docutils literal notranslate"><span class="pre">barrier_wait</span></code></a>(barrier_id[, parity])</p></td>
+<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.builtin.barrier_wait" title="tilelang.language.builtin.barrier_wait"><code class="xref py py-obj docutils literal notranslate"><span class="pre">barrier_wait</span></code></a>(barrier_id[, parity])</p></td>
 <td><p>Wait for a memory barrier to complete.</p></td>
 </tr>
-<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.builtin.barrier_arrive" title="tilelang.language.builtin.barrier_arrive"><code class="xref py py-obj docutils literal notranslate"><span class="pre">barrier_arrive</span></code></a>(barrier_id)</p></td>
+<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.builtin.barrier_arrive" title="tilelang.language.builtin.barrier_arrive"><code class="xref py py-obj docutils literal notranslate"><span class="pre">barrier_arrive</span></code></a>(barrier_id)</p></td>
 <td><p>Arrive at a memory barrier.</p></td>
 </tr>
-<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.builtin.shfl_xor" title="tilelang.language.builtin.shfl_xor"><code class="xref py py-obj docutils literal notranslate"><span class="pre">shfl_xor</span></code></a>(value, offset)</p></td>
+<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.builtin.shfl_xor" title="tilelang.language.builtin.shfl_xor"><code class="xref py py-obj docutils literal notranslate"><span class="pre">shfl_xor</span></code></a>(value, offset)</p></td>
 <td><p>Perform a shuffle operation with XOR offset.</p></td>
 </tr>
-<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.builtin.shfl_down" title="tilelang.language.builtin.shfl_down"><code class="xref py py-obj docutils literal notranslate"><span class="pre">shfl_down</span></code></a>(value, offset)</p></td>
+<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.builtin.shfl_down" title="tilelang.language.builtin.shfl_down"><code class="xref py py-obj docutils literal notranslate"><span class="pre">shfl_down</span></code></a>(value, offset)</p></td>
 <td><p>Perform a shuffle operation with down offset.</p></td>
 </tr>
-<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.builtin.shfl_up" title="tilelang.language.builtin.shfl_up"><code class="xref py py-obj docutils literal notranslate"><span class="pre">shfl_up</span></code></a>(value, offset)</p></td>
+<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.builtin.shfl_up" title="tilelang.language.builtin.shfl_up"><code class="xref py py-obj docutils literal notranslate"><span class="pre">shfl_up</span></code></a>(value, offset)</p></td>
 <td><p>Perform a shuffle operation with up offset.</p></td>
 </tr>
-<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.builtin.sync_threads" title="tilelang.language.builtin.sync_threads"><code class="xref py py-obj docutils literal notranslate"><span class="pre">sync_threads</span></code></a>()</p></td>
+<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.builtin.sync_threads" title="tilelang.language.builtin.sync_threads"><code class="xref py py-obj docutils literal notranslate"><span class="pre">sync_threads</span></code></a>()</p></td>
 <td><p>Synchronize all threads in a warp.</p></td>
 </tr>
-<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.builtin.sync_thread_partial" title="tilelang.language.builtin.sync_thread_partial"><code class="xref py py-obj docutils literal notranslate"><span class="pre">sync_thread_partial</span></code></a>(barrier_id)</p></td>
+<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.builtin.sync_thread_partial" title="tilelang.language.builtin.sync_thread_partial"><code class="xref py py-obj docutils literal notranslate"><span class="pre">sync_thread_partial</span></code></a>(barrier_id)</p></td>
 <td><p>Synchronize threads within a warp.</p></td>
 </tr>
-<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.builtin.sync_global" title="tilelang.language.builtin.sync_global"><code class="xref py py-obj docutils literal notranslate"><span class="pre">sync_global</span></code></a>()</p></td>
+<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.builtin.sync_global" title="tilelang.language.builtin.sync_global"><code class="xref py py-obj docutils literal notranslate"><span class="pre">sync_global</span></code></a>()</p></td>
 <td><p>Synchronize all threads in a block.</p></td>
 </tr>
-<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.language.builtin.sync_grid" title="tilelang.language.builtin.sync_grid"><code class="xref py py-obj docutils literal notranslate"><span class="pre">sync_grid</span></code></a>()</p></td>
+<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.language.builtin.sync_grid" title="tilelang.language.builtin.sync_grid"><code class="xref py py-obj docutils literal notranslate"><span class="pre">sync_grid</span></code></a>()</p></td>
 <td><p>Synchronize all threads in a grid.</p></td>
 </tr>
 </tbody>
@@ -718,12 +727,40 @@ <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.builtin.annotate_producer_reg_dealloc">
+<span class="sig-prename descclassname"><span class="pre">tilelang.language.builtin.</span></span><span class="sig-name descname"><span class="pre">annotate_producer_reg_dealloc</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">reg_count</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">24</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#tilelang.language.builtin.annotate_producer_reg_dealloc" title="Link to this definition">¶</a></dt>
+<dd><p>Annotate the producer reg dealloc.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters<span class="colon">:</span></dt>
+<dd class="field-odd"><p><strong>reg_count</strong> (<em>int</em>)</p>
+</dd>
+</dl>
+</dd></dl>
+
+<dl class="py function">
+<dt class="sig sig-object py" id="tilelang.language.builtin.annotate_consumer_reg_alloc">
+<span class="sig-prename descclassname"><span class="pre">tilelang.language.builtin.</span></span><span class="sig-name descname"><span class="pre">annotate_consumer_reg_alloc</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">reg_count</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">240</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#tilelang.language.builtin.annotate_consumer_reg_alloc" title="Link to this definition">¶</a></dt>
+<dd><p>Annotate the consumer reg alloc.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters<span class="colon">:</span></dt>
+<dd class="field-odd"><p><strong>reg_count</strong> (<em>int</em>)</p>
+</dd>
+</dl>
+</dd></dl>
+
 <dl class="py function">
 <dt class="sig sig-object py" id="tilelang.language.builtin.no_set_max_nreg">
 <span class="sig-prename descclassname"><span class="pre">tilelang.language.builtin.</span></span><span class="sig-name descname"><span class="pre">no_set_max_nreg</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="headerlink" href="#tilelang.language.builtin.no_set_max_nreg" title="Link to this definition">¶</a></dt>
 <dd><p>Disable the maximum register limit setting.</p>
 </dd></dl>
 
+<dl class="py function">
+<dt class="sig sig-object py" id="tilelang.language.builtin.disable_warp_group_reg_alloc">
+<span class="sig-prename descclassname"><span class="pre">tilelang.language.builtin.</span></span><span class="sig-name descname"><span class="pre">disable_warp_group_reg_alloc</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="headerlink" href="#tilelang.language.builtin.disable_warp_group_reg_alloc" title="Link to this definition">¶</a></dt>
+<dd><p>Disable the warp group reg alloc.</p>
+</dd></dl>
+
 <dl class="py function">
 <dt class="sig sig-object py" id="tilelang.language.builtin.mbarrier_wait_parity">
 <span class="sig-prename descclassname"><span class="pre">tilelang.language.builtin.</span></span><span class="sig-name descname"><span class="pre">mbarrier_wait_parity</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">mbarrier</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">parity</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#tilelang.language.builtin.mbarrier_wait_parity" title="Link to this definition">¶</a></dt>
@@ -1023,7 +1060,10 @@ <h2>Module Contents<a class="headerlink" href="#module-contents" title="Link to
 <li><a class="reference internal" href="#tilelang.language.builtin.set_max_nreg"><code class="docutils literal notranslate"><span class="pre">set_max_nreg()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.language.builtin.inc_max_nreg"><code class="docutils literal notranslate"><span class="pre">inc_max_nreg()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.language.builtin.dec_max_nreg"><code class="docutils literal notranslate"><span class="pre">dec_max_nreg()</span></code></a></li>
+<li><a class="reference internal" href="#tilelang.language.builtin.annotate_producer_reg_dealloc"><code class="docutils literal notranslate"><span class="pre">annotate_producer_reg_dealloc()</span></code></a></li>
+<li><a class="reference internal" href="#tilelang.language.builtin.annotate_consumer_reg_alloc"><code class="docutils literal notranslate"><span class="pre">annotate_consumer_reg_alloc()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.language.builtin.no_set_max_nreg"><code class="docutils literal notranslate"><span class="pre">no_set_max_nreg()</span></code></a></li>
+<li><a class="reference internal" href="#tilelang.language.builtin.disable_warp_group_reg_alloc"><code class="docutils literal notranslate"><span class="pre">disable_warp_group_reg_alloc()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.language.builtin.mbarrier_wait_parity"><code class="docutils literal notranslate"><span class="pre">mbarrier_wait_parity()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.language.builtin.mbarrier_arrive"><code class="docutils literal notranslate"><span class="pre">mbarrier_arrive()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.language.builtin.mbarrier_expect_tx"><code class="docutils literal notranslate"><span class="pre">mbarrier_expect_tx()</span></code></a></li>