Update docs

Author: github-actions[bot]

_sources/autoapi/tilelang/intrinsics/tcgen05_macro_generator/index.rst.txt

@@ -120,6 +120,25 @@ Module Contents
 
 
 
+   .. py:method:: tcgen05mma_ss(A_buf, B_buf, C_local_buf, mbar, clear_accum = False)
+
+      Emit the SS (Shared-Shared) variant of TCGEN5MMA.
+
+      Reads operand A and B from shared memory via a descriptor.
+
+      :param A_buf: Operand A in shared memory.
+      :type A_buf: Buffer
+      :param B_buf: Operand B in shared memory.
+      :type B_buf: Buffer
+      :param C_local_buf: Accumulator buffer in tensor memory.
+      :type C_local_buf: Buffer
+      :param mbar: Memory barrier for MMA completion signalling.
+      :type mbar: PrimExpr
+      :param clear_accum: Whether to zero the accumulator before the first MMA.
+      :type clear_accum: PrimExpr
+
+
+
    .. py:method:: tcgen05mma_ts(A_buf, B_buf, C_local_buf, mbar, clear_accum = False)
 
       Emit the TS (TensorMemory-Shared) variant of TCGEN5MMA.
@@ -176,9 +195,9 @@ Module Contents
 
 
 
-   .. py:method:: get_tcgen5_mma_meta(m, n, k)
+   .. py:method:: get_tcgen5_mma_meta(m, n, k, disable_2cta)
 
-      Query the FFI for TCGEN5MMA atom metadata (atom_m, atom_n, atom_k, enable_ws, enable_2cta).
+      Query the FFI for TCGEN5MMA atom metadata (atom_m, atom_n, atom_k, enable_ws, enable_2cta), and record them in `self.meta`.
 
 
 

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

@@ -17,6 +17,7 @@ Functions
    tilelang.language.builtin.access_ptr
    tilelang.language.builtin.create_tma_descriptor
    tilelang.language.builtin.tma_load
+   tilelang.language.builtin.tma_load_2sm
    tilelang.language.builtin.fence_proxy_async
    tilelang.language.builtin.tma_store_arrive
    tilelang.language.builtin.tma_store_wait
@@ -135,6 +136,16 @@ Module Contents
    :rtype: tir.Call
 
 
+.. py:function:: tma_load_2sm(*args)
+
+   Perform a Tensor Memory Access (TMA) load operation with 2SM on Blackwell.
+
+   :param \*args: Variable arguments specifying the TMA load parameters
+
+   :returns: A handle to the TMA load operation
+   :rtype: tir.Call
+
+
 .. py:function:: fence_proxy_async(*args)
 
    Create a fence for asynchronous proxy operations.
@@ -556,12 +567,15 @@ Module Contents
    Perform a ptx async copy barrier using cp.async.mbarrier.arrive.noinc.
 
 
-.. py:function:: tcgen05_mma_arrive(mbar)
+.. py:function:: tcgen05_mma_arrive(mbar, arrive_2cta = False)
 
    Signal UMMA (TCGEN05) barrier arrival for a shared-memory mbarrier pointer.
 
    :param mbar: The mbarrier object in shared memory (e.g., Barrier*) or its address.
    :type mbar: tir.Buffer | BufferLoad | PrimExpr
+   :param arrive_2cta: Whether to also arrive at the peer CTA's barrier.
+                       If set, will be lowered to umma_arrive_multicast_2x1SM.
+   :type arrive_2cta: bool
 
 
 .. py:function:: ptx_mma_sm70(shape, A_layout, B_layout, A_dtype, B_dtype, C_dtype, multiplicand_a, a_index, multiplicand_b, b_index, accumulator, c_index)

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

@@ -84,7 +84,7 @@ Module Contents
    compilation fails instead of silently falling back to MMA.
 
 
-.. py:function:: tcgen05_gemm(A, B, C, transpose_A = False, transpose_B = False, policy = GemmWarpPolicy.Square, clear_accum = False, *, mbar)
+.. py:function:: tcgen05_gemm(A, B, C, transpose_A = False, transpose_B = False, policy = GemmWarpPolicy.Square, clear_accum = False, *, mbar, use_2cta = False)
 
    Explicit Blackwell TCGEN05 GEMM without an implicit wait.
 
@@ -93,6 +93,9 @@ Module Contents
    - it always requests the TCGEN5MMA lowering path
    - it never auto-emits an inlined `mbarrier_wait_parity`
 
+   When ``use_2cta=True``, the instruction is lowered to the 2CTA variant
+   which requires ``cluster_dims`` to be ``(2,1,1)`` or ``(1,2,1)``.
+
    If the current target or operand pattern cannot use Blackwell TCGEN5MMA,
    compilation fails instead of silently falling back to another GEMM path.
 

_sources/autoapi/tilelang/language/tir/op/index.rst.txt

@@ -872,11 +872,11 @@ Module Contents
 
 .. py:function:: ptx_wgmma_rs(dtype, wgmma_prefix, b_is_k_major, a_dtype_abbrv, b_dtype_abbrv, accum_dtype_abbrv, A_buf, A_offset, B_desc, B_offset, C_data, C_offset, scale_out, scale_in_a, scale_in_b)
 
-.. py:function:: ptx_tcgen05_mma_ss(kind_dtype, desc_a, A_offset, desc_b, B_offset, C_ptr, C_offset, desc_val, scale_out, mask0, mask1, mask2, mask3, enable_ws=False, ws=None, warp_specialized=None, variant=None)
+.. py:function:: ptx_tcgen05_mma_ss(kind_dtype, desc_a, A_offset, desc_b, B_offset, C_ptr, C_offset, desc_val, scale_out, mask0, mask1, mask2, mask3, enable_ws=False, enable_2cta=False, ws=None, warp_specialized=None, variant=None)
 
-   TVM intrinsic for tcgen05.mma shared-memory × shared-memory instructions.
+   TVM intrinsic for tcgen05.mma shared-memory x shared-memory instructions.
 
-   Expects 13 or 14 positional arguments:
+   Expects 14 or 15 positional arguments:
    (kind_dtype, desc_a, A_offset, desc_b, B_offset, C_ptr, C_offset,
     desc_val, scale_out, mask0, mask1, mask2, mask3[, enable_ws]).
    Aliases: you can also pass `ws` or `warp_specialized` (booleans) instead of `enable_ws`.
@@ -885,9 +885,9 @@ Module Contents
      "tf32" for kind::tf32, "int8" for kind::i8, "float8_e4m3" for kind::f8f6f4).
 
 
-.. py:function:: ptx_tcgen05_mma_ts(kind_dtype, A_ptr, A_offset, desc_b, B_offset, C_ptr, C_offset, desc_val, scale_out, mask0, mask1, mask2, mask3)
+.. py:function:: ptx_tcgen05_mma_ts(kind_dtype, A_ptr, A_offset, desc_b, B_offset, C_ptr, C_offset, desc_val, scale_out, mask0, mask1, mask2, mask3, enable_2cta=False)
 
-   TVM intrinsic for tcgen05.mma tensor-memory × shared-memory instructions.
+   TVM intrinsic for tcgen05.mma tensor-memory x shared-memory instructions.
 
    Expects 13 positional arguments:
    (kind_dtype, A_ptr, A_offset, desc_b, B_offset, C_ptr, C_offset,

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

@@ -84,6 +84,7 @@ Functions
    tilelang.transform.LayoutReducer
    tilelang.transform.UnrollLoop
    tilelang.transform.LowerLDGSTG
+   tilelang.transform.LowerBlackwell2SM
 
 
 Package Contents
@@ -547,3 +548,14 @@ Package Contents
    :rtype: tvm.transform.Pass
 
 
+.. py:function:: LowerBlackwell2SM()
+
+   Lower 2SM TCGEN5MMA and related on Blackwell target
+
+   :returns:
+
+             tvm.transform.Pass
+                 The result pass
+   :rtype: fpass
+
+

autoapi/tilelang/intrinsics/tcgen05_macro_generator/index.html

@@ -745,6 +745,24 @@ <h2>Module Contents<a class="headerlink" href="#module-contents" title="Link to
 </dl>
 </dd></dl>
 
+<dl class="py method">
+<dt class="sig sig-object py" id="tilelang.intrinsics.tcgen05_macro_generator.TensorCoreIntrinEmitter.tcgen05mma_ss">
+<span class="sig-name descname"><span class="pre">tcgen05mma_ss</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">A_buf</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">B_buf</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">C_local_buf</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">mbar</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">clear_accum</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.intrinsics.tcgen05_macro_generator.TensorCoreIntrinEmitter.tcgen05mma_ss" title="Link to this definition">¶</a></dt>
+<dd><p>Emit the SS (Shared-Shared) variant of TCGEN5MMA.</p>
+<p>Reads operand A and B from shared memory via a descriptor.</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>A_buf</strong> (<em>Buffer</em>) – Operand A in shared memory.</p></li>
+<li><p><strong>B_buf</strong> (<em>Buffer</em>) – Operand B in shared memory.</p></li>
+<li><p><strong>C_local_buf</strong> (<em>Buffer</em>) – Accumulator buffer in tensor memory.</p></li>
+<li><p><strong>mbar</strong> (<em>PrimExpr</em>) – Memory barrier for MMA completion signalling.</p></li>
+<li><p><strong>clear_accum</strong> (<em>PrimExpr</em>) – Whether to zero the accumulator before the first MMA.</p></li>
+</ul>
+</dd>
+</dl>
+</dd></dl>
+
 <dl class="py method">
 <dt class="sig sig-object py" id="tilelang.intrinsics.tcgen05_macro_generator.TensorCoreIntrinEmitter.tcgen05mma_ts">
 <span class="sig-name descname"><span class="pre">tcgen05mma_ts</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">A_buf</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">B_buf</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">C_local_buf</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">mbar</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">clear_accum</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.intrinsics.tcgen05_macro_generator.TensorCoreIntrinEmitter.tcgen05mma_ts" title="Link to this definition">¶</a></dt>
@@ -815,14 +833,15 @@ <h2>Module Contents<a class="headerlink" href="#module-contents" title="Link to
 
 <dl class="py method">
 <dt class="sig sig-object py" id="tilelang.intrinsics.tcgen05_macro_generator.TensorCoreIntrinEmitter.get_tcgen5_mma_meta">
-<span class="sig-name descname"><span class="pre">get_tcgen5_mma_meta</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">m</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">n</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">k</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#tilelang.intrinsics.tcgen05_macro_generator.TensorCoreIntrinEmitter.get_tcgen5_mma_meta" title="Link to this definition">¶</a></dt>
-<dd><p>Query the FFI for TCGEN5MMA atom metadata (atom_m, atom_n, atom_k, enable_ws, enable_2cta).</p>
+<span class="sig-name descname"><span class="pre">get_tcgen5_mma_meta</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">m</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">n</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">k</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">disable_2cta</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#tilelang.intrinsics.tcgen05_macro_generator.TensorCoreIntrinEmitter.get_tcgen5_mma_meta" title="Link to this definition">¶</a></dt>
+<dd><p>Query the FFI for TCGEN5MMA atom metadata (atom_m, atom_n, atom_k, enable_ws, enable_2cta), and record them in <cite>self.meta</cite>.</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>m</strong> (<em>int</em>)</p></li>
 <li><p><strong>n</strong> (<em>int</em>)</p></li>
 <li><p><strong>k</strong> (<em>int</em>)</p></li>
+<li><p><strong>disable_2cta</strong> (<a class="reference internal" href="../../language/dtypes/index.html#tilelang.language.dtypes.bool" title="tilelang.language.dtypes.bool"><em>bool</em></a>)</p></li>
 </ul>
 </dd>
 </dl>
@@ -933,6 +952,7 @@ <h2>Module Contents<a class="headerlink" href="#module-contents" title="Link to
 <li><a class="reference internal" href="#tilelang.intrinsics.tcgen05_macro_generator.TensorCoreIntrinEmitter.a_shared_layout"><code class="docutils literal notranslate"><span class="pre">TensorCoreIntrinEmitter.a_shared_layout</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.intrinsics.tcgen05_macro_generator.TensorCoreIntrinEmitter.b_shared_layout"><code class="docutils literal notranslate"><span class="pre">TensorCoreIntrinEmitter.b_shared_layout</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.intrinsics.tcgen05_macro_generator.TensorCoreIntrinEmitter.tcgen05mma"><code class="docutils literal notranslate"><span class="pre">TensorCoreIntrinEmitter.tcgen05mma()</span></code></a></li>
+<li><a class="reference internal" href="#tilelang.intrinsics.tcgen05_macro_generator.TensorCoreIntrinEmitter.tcgen05mma_ss"><code class="docutils literal notranslate"><span class="pre">TensorCoreIntrinEmitter.tcgen05mma_ss()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.intrinsics.tcgen05_macro_generator.TensorCoreIntrinEmitter.tcgen05mma_ts"><code class="docutils literal notranslate"><span class="pre">TensorCoreIntrinEmitter.tcgen05mma_ts()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.intrinsics.tcgen05_macro_generator.TensorCoreIntrinEmitter.make_mma_load_layout"><code class="docutils literal notranslate"><span class="pre">TensorCoreIntrinEmitter.make_mma_load_layout()</span></code></a></li>
 <li><a class="reference internal" href="#tilelang.intrinsics.tcgen05_macro_generator.TensorCoreIntrinEmitter.make_mma_store_layout"><code class="docutils literal notranslate"><span class="pre">TensorCoreIntrinEmitter.make_mma_store_layout()</span></code></a></li>