Skip to content

Commit 4a82073

Browse files
Update docs
1 parent fb625e1 commit 4a82073

File tree

7 files changed

+298
-54
lines changed

7 files changed

+298
-54
lines changed

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

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,9 @@ Attributes
1414
tilelang.env.COMPOSABLE_KERNEL_NOT_FOUND_MESSAGE
1515
tilelang.env.TL_TEMPLATE_NOT_FOUND_MESSAGE
1616
tilelang.env.TVM_LIBRARY_NOT_FOUND_MESSAGE
17-
tilelang.env.SITE_PACKAGES
17+
tilelang.env.TL_ROOT
1818
tilelang.env.TL_LIBS
1919
tilelang.env.TL_LIBS
20-
tilelang.env.TL_ROOT
2120
tilelang.env.DEV
2221
tilelang.env.THIRD_PARTY_ROOT
2322
tilelang.env.DEV
@@ -81,16 +80,12 @@ Module Contents
8180
:value: 'TVM is not installed or found in the expected path'
8281

8382

84-
.. py:data:: SITE_PACKAGES
85-
:value: []
86-
83+
.. py:data:: TL_ROOT
8784
8885
.. py:data:: TL_LIBS
8986
9087
.. py:data:: TL_LIBS
9188
92-
.. py:data:: TL_ROOT
93-
9489
.. py:data:: DEV
9590
:value: False
9691

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

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,11 @@ Functions
3434
tilelang.language.builtin.warpgroup_arrive
3535
tilelang.language.builtin.warpgroup_commit_batch
3636
tilelang.language.builtin.warpgroup_wait
37+
tilelang.language.builtin.get_lane_idx
38+
tilelang.language.builtin.get_warp_idx_sync
39+
tilelang.language.builtin.get_warp_idx
40+
tilelang.language.builtin.get_warp_group_idx
41+
tilelang.language.builtin.shuffle_elect
3742
tilelang.language.builtin.wait_wgmma
3843
tilelang.language.builtin.barrier_wait
3944
tilelang.language.builtin.barrier_arrive
@@ -270,6 +275,103 @@ Module Contents
270275
:rtype: tir.Call
271276

272277

278+
.. py:function:: get_lane_idx(warp_size = None)
279+
280+
Return the logical lane index of the calling thread within a warp.
281+
282+
:param warp_size: Logical warp (or wavefront) size. Defaults to 32 on NVIDIA and 64 on AMD.
283+
:type warp_size: Optional[int, PrimExpr]
284+
285+
.. rubric:: Example
286+
287+
>>> lane = T.get_lane_idx()
288+
>>> custom_lane = T.get_lane_idx(64) # override warp size explicitly
289+
290+
Implementation Notes
291+
--------------------
292+
Lowers to the CUDA helper `tl::get_lane_idx(warp_size)` defined in
293+
`src/tl_templates/cuda/intrin.h`, which computes the lane index from the
294+
linear thread id using the provided `warp_size`.
295+
296+
297+
.. py:function:: get_warp_idx_sync(warp_size = None)
298+
299+
Return the canonical warp index, assuming the warp's threads are converged.
300+
301+
:param warp_size: Logical warp size used for the index calculation.
302+
:type warp_size: Optional[int, PrimExpr]
303+
304+
.. rubric:: Example
305+
306+
>>> warp = T.get_warp_idx_sync()
307+
>>> custom_warp = T.get_warp_idx_sync(64)
308+
309+
Implementation Notes
310+
--------------------
311+
Emits `tl::get_warp_idx_sync(warp_size)` which divides the block-linear
312+
thread id by `warp_size`, matching the semantics of CUTLASS' canonical helpers.
313+
314+
315+
.. py:function:: get_warp_idx(warp_size = None)
316+
317+
Return the canonical warp index without synchronizing the warp.
318+
319+
:param warp_size: Logical warp size used for the index calculation.
320+
:type warp_size: Optional[int, PrimExpr]
321+
322+
.. rubric:: Example
323+
324+
>>> warp = T.get_warp_idx()
325+
>>> custom_warp = T.get_warp_idx(64)
326+
327+
Implementation Notes
328+
--------------------
329+
Lowers to `tl::get_warp_idx(warp_size)` which divides the block-linear
330+
thread id by the provided `warp_size` without requiring warp convergence.
331+
332+
333+
.. py:function:: get_warp_group_idx(warp_size = None, warps_per_group = None)
334+
335+
Return the canonical warp group index for the calling thread.
336+
337+
:param warp_size: Logical warp size to use (defaults to 32 on NVIDIA / 64 on AMD).
338+
:type warp_size: Optional[int, PrimExpr]
339+
:param warps_per_group: Number of warps per warp-group. Defaults to 4 on NVIDIA architectures.
340+
:type warps_per_group: Optional[int, PrimExpr]
341+
342+
.. rubric:: Example
343+
344+
>>> group = T.get_warp_group_idx()
345+
>>> custom_group = T.get_warp_group_idx(32, 6) # treat 6 warps as a group
346+
347+
Implementation Notes
348+
--------------------
349+
Generates `tl::get_warp_group_idx(warp_size, warps_per_group)` which
350+
divides the block-linear thread id by `warp_size * warps_per_group`,
351+
matching the canonical ordering while allowing architecture-specific overrides.
352+
353+
354+
.. py:function:: shuffle_elect(thread_extent)
355+
356+
Elect exactly one lane within a logical thread group.
357+
358+
:param thread_extent: Size (in threads) of the group in which a single lane should be elected.
359+
Passing 0 elects a single lane in the entire thread block.
360+
:type thread_extent: int
361+
362+
.. rubric:: Example
363+
364+
>>> is_leader = T.shuffle_elect(64)
365+
>>> T.if_then_else(is_leader, do_leader_work(), T.evaluate(0))
366+
367+
Implementation Notes
368+
--------------------
369+
Lowered to the CUDA helper `tl::tl_shuffle_elect<thread_extent>()` defined in
370+
`src/tl_templates/cuda/intrin.h`, which relies on
371+
`cutlass::canonical_warp_idx_sync()` and `cute::elect_one_sync()` (or
372+
`__shfl_sync`) to pick one lane per group.
373+
374+
273375
.. py:function:: wait_wgmma(id)
274376
275377
Wait for WGMMA (Warp Group Matrix Multiply-Accumulate) operations to complete.

autoapi/tilelang/env/index.html

Lines changed: 20 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -498,7 +498,7 @@ <h2>Attributes<a class="headerlink" href="#attributes" title="Link to this headi
498498
<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.env.TVM_LIBRARY_NOT_FOUND_MESSAGE" title="tilelang.env.TVM_LIBRARY_NOT_FOUND_MESSAGE"><code class="xref py py-obj docutils literal notranslate"><span class="pre">TVM_LIBRARY_NOT_FOUND_MESSAGE</span></code></a></p></td>
499499
<td><p></p></td>
500500
</tr>
501-
<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.env.SITE_PACKAGES" title="tilelang.env.SITE_PACKAGES"><code class="xref py py-obj docutils literal notranslate"><span class="pre">SITE_PACKAGES</span></code></a></p></td>
501+
<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.env.TL_ROOT" title="tilelang.env.TL_ROOT"><code class="xref py py-obj docutils literal notranslate"><span class="pre">TL_ROOT</span></code></a></p></td>
502502
<td><p></p></td>
503503
</tr>
504504
<tr class="row-odd"><td><p><a class="reference internal" href="#id0" title="tilelang.env.TL_LIBS"><code class="xref py py-obj docutils literal notranslate"><span class="pre">TL_LIBS</span></code></a></p></td>
@@ -507,55 +507,52 @@ <h2>Attributes<a class="headerlink" href="#attributes" title="Link to this headi
507507
<tr class="row-even"><td><p><a class="reference internal" href="#id0" title="tilelang.env.TL_LIBS"><code class="xref py py-obj docutils literal notranslate"><span class="pre">TL_LIBS</span></code></a></p></td>
508508
<td><p></p></td>
509509
</tr>
510-
<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.env.TL_ROOT" title="tilelang.env.TL_ROOT"><code class="xref py py-obj docutils literal notranslate"><span class="pre">TL_ROOT</span></code></a></p></td>
510+
<tr class="row-odd"><td><p><a class="reference internal" href="#id1" title="tilelang.env.DEV"><code class="xref py py-obj docutils literal notranslate"><span class="pre">DEV</span></code></a></p></td>
511511
<td><p></p></td>
512512
</tr>
513-
<tr class="row-even"><td><p><a class="reference internal" href="#id1" title="tilelang.env.DEV"><code class="xref py py-obj docutils literal notranslate"><span class="pre">DEV</span></code></a></p></td>
513+
<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.env.THIRD_PARTY_ROOT" title="tilelang.env.THIRD_PARTY_ROOT"><code class="xref py py-obj docutils literal notranslate"><span class="pre">THIRD_PARTY_ROOT</span></code></a></p></td>
514514
<td><p></p></td>
515515
</tr>
516-
<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.env.THIRD_PARTY_ROOT" title="tilelang.env.THIRD_PARTY_ROOT"><code class="xref py py-obj docutils literal notranslate"><span class="pre">THIRD_PARTY_ROOT</span></code></a></p></td>
516+
<tr class="row-odd"><td><p><a class="reference internal" href="#id1" title="tilelang.env.DEV"><code class="xref py py-obj docutils literal notranslate"><span class="pre">DEV</span></code></a></p></td>
517517
<td><p></p></td>
518518
</tr>
519-
<tr class="row-even"><td><p><a class="reference internal" href="#id1" title="tilelang.env.DEV"><code class="xref py py-obj docutils literal notranslate"><span class="pre">DEV</span></code></a></p></td>
519+
<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.env.enable_cache" title="tilelang.env.enable_cache"><code class="xref py py-obj docutils literal notranslate"><span class="pre">enable_cache</span></code></a></p></td>
520520
<td><p></p></td>
521521
</tr>
522-
<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.env.enable_cache" title="tilelang.env.enable_cache"><code class="xref py py-obj docutils literal notranslate"><span class="pre">enable_cache</span></code></a></p></td>
522+
<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.env.disable_cache" title="tilelang.env.disable_cache"><code class="xref py py-obj docutils literal notranslate"><span class="pre">disable_cache</span></code></a></p></td>
523523
<td><p></p></td>
524524
</tr>
525-
<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.env.disable_cache" title="tilelang.env.disable_cache"><code class="xref py py-obj docutils literal notranslate"><span class="pre">disable_cache</span></code></a></p></td>
525+
<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.env.is_cache_enabled" title="tilelang.env.is_cache_enabled"><code class="xref py py-obj docutils literal notranslate"><span class="pre">is_cache_enabled</span></code></a></p></td>
526526
<td><p></p></td>
527527
</tr>
528-
<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.env.is_cache_enabled" title="tilelang.env.is_cache_enabled"><code class="xref py py-obj docutils literal notranslate"><span class="pre">is_cache_enabled</span></code></a></p></td>
528+
<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.env.env" title="tilelang.env.env"><code class="xref py py-obj docutils literal notranslate"><span class="pre">env</span></code></a></p></td>
529529
<td><p></p></td>
530530
</tr>
531-
<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.env.env" title="tilelang.env.env"><code class="xref py py-obj docutils literal notranslate"><span class="pre">env</span></code></a></p></td>
531+
<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.env.CUDA_HOME" title="tilelang.env.CUDA_HOME"><code class="xref py py-obj docutils literal notranslate"><span class="pre">CUDA_HOME</span></code></a></p></td>
532532
<td><p></p></td>
533533
</tr>
534-
<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.env.CUDA_HOME" title="tilelang.env.CUDA_HOME"><code class="xref py py-obj docutils literal notranslate"><span class="pre">CUDA_HOME</span></code></a></p></td>
534+
<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.env.ROCM_HOME" title="tilelang.env.ROCM_HOME"><code class="xref py py-obj docutils literal notranslate"><span class="pre">ROCM_HOME</span></code></a></p></td>
535535
<td><p></p></td>
536536
</tr>
537-
<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.env.ROCM_HOME" title="tilelang.env.ROCM_HOME"><code class="xref py py-obj docutils literal notranslate"><span class="pre">ROCM_HOME</span></code></a></p></td>
537+
<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.env.tvm_path" title="tilelang.env.tvm_path"><code class="xref py py-obj docutils literal notranslate"><span class="pre">tvm_path</span></code></a></p></td>
538538
<td><p></p></td>
539539
</tr>
540-
<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.env.tvm_path" title="tilelang.env.tvm_path"><code class="xref py py-obj docutils literal notranslate"><span class="pre">tvm_path</span></code></a></p></td>
540+
<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.env.cutlass_inc_path" title="tilelang.env.cutlass_inc_path"><code class="xref py py-obj docutils literal notranslate"><span class="pre">cutlass_inc_path</span></code></a></p></td>
541541
<td><p></p></td>
542542
</tr>
543-
<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.env.cutlass_inc_path" title="tilelang.env.cutlass_inc_path"><code class="xref py py-obj docutils literal notranslate"><span class="pre">cutlass_inc_path</span></code></a></p></td>
543+
<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.env.ck_inc_path" title="tilelang.env.ck_inc_path"><code class="xref py py-obj docutils literal notranslate"><span class="pre">ck_inc_path</span></code></a></p></td>
544544
<td><p></p></td>
545545
</tr>
546-
<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.env.ck_inc_path" title="tilelang.env.ck_inc_path"><code class="xref py py-obj docutils literal notranslate"><span class="pre">ck_inc_path</span></code></a></p></td>
546+
<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.env.tl_template_path" title="tilelang.env.tl_template_path"><code class="xref py py-obj docutils literal notranslate"><span class="pre">tl_template_path</span></code></a></p></td>
547547
<td><p></p></td>
548548
</tr>
549-
<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.env.tl_template_path" title="tilelang.env.tl_template_path"><code class="xref py py-obj docutils literal notranslate"><span class="pre">tl_template_path</span></code></a></p></td>
549+
<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.env.CUTLASS_INCLUDE_DIR" title="tilelang.env.CUTLASS_INCLUDE_DIR"><code class="xref py py-obj docutils literal notranslate"><span class="pre">CUTLASS_INCLUDE_DIR</span></code></a></p></td>
550550
<td><p></p></td>
551551
</tr>
552-
<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.env.CUTLASS_INCLUDE_DIR" title="tilelang.env.CUTLASS_INCLUDE_DIR"><code class="xref py py-obj docutils literal notranslate"><span class="pre">CUTLASS_INCLUDE_DIR</span></code></a></p></td>
552+
<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.env.COMPOSABLE_KERNEL_INCLUDE_DIR" title="tilelang.env.COMPOSABLE_KERNEL_INCLUDE_DIR"><code class="xref py py-obj docutils literal notranslate"><span class="pre">COMPOSABLE_KERNEL_INCLUDE_DIR</span></code></a></p></td>
553553
<td><p></p></td>
554554
</tr>
555-
<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.env.COMPOSABLE_KERNEL_INCLUDE_DIR" title="tilelang.env.COMPOSABLE_KERNEL_INCLUDE_DIR"><code class="xref py py-obj docutils literal notranslate"><span class="pre">COMPOSABLE_KERNEL_INCLUDE_DIR</span></code></a></p></td>
556-
<td><p></p></td>
557-
</tr>
558-
<tr class="row-odd"><td><p><a class="reference internal" href="#tilelang.env.TILELANG_TEMPLATE_PATH" title="tilelang.env.TILELANG_TEMPLATE_PATH"><code class="xref py py-obj docutils literal notranslate"><span class="pre">TILELANG_TEMPLATE_PATH</span></code></a></p></td>
555+
<tr class="row-even"><td><p><a class="reference internal" href="#tilelang.env.TILELANG_TEMPLATE_PATH" title="tilelang.env.TILELANG_TEMPLATE_PATH"><code class="xref py py-obj docutils literal notranslate"><span class="pre">TILELANG_TEMPLATE_PATH</span></code></a></p></td>
559556
<td><p></p></td>
560557
</tr>
561558
</tbody>
@@ -623,8 +620,8 @@ <h2>Module Contents<a class="headerlink" href="#module-contents" title="Link to
623620
<dd></dd></dl>
624621

625622
<dl class="py data">
626-
<dt class="sig sig-object py" id="tilelang.env.SITE_PACKAGES">
627-
<span class="sig-prename descclassname"><span class="pre">tilelang.env.</span></span><span class="sig-name descname"><span class="pre">SITE_PACKAGES</span></span><em class="property"><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><span class="pre">[]</span></em><a class="headerlink" href="#tilelang.env.SITE_PACKAGES" title="Link to this definition"></a></dt>
623+
<dt class="sig sig-object py" id="tilelang.env.TL_ROOT">
624+
<span class="sig-prename descclassname"><span class="pre">tilelang.env.</span></span><span class="sig-name descname"><span class="pre">TL_ROOT</span></span><a class="headerlink" href="#tilelang.env.TL_ROOT" title="Link to this definition"></a></dt>
628625
<dd></dd></dl>
629626

630627
<dl class="py data">
@@ -637,11 +634,6 @@ <h2>Module Contents<a class="headerlink" href="#module-contents" title="Link to
637634
<span class="sig-prename descclassname"><span class="pre">tilelang.env.</span></span><span class="sig-name descname"><span class="pre">TL_LIBS</span></span><a class="headerlink" href="#id0" title="Link to this definition"></a></dt>
638635
<dd></dd></dl>
639636

640-
<dl class="py data">
641-
<dt class="sig sig-object py" id="tilelang.env.TL_ROOT">
642-
<span class="sig-prename descclassname"><span class="pre">tilelang.env.</span></span><span class="sig-name descname"><span class="pre">TL_ROOT</span></span><a class="headerlink" href="#tilelang.env.TL_ROOT" title="Link to this definition"></a></dt>
643-
<dd></dd></dl>
644-
645637
<dl class="py data">
646638
<dt class="sig sig-object py" id="tilelang.env.DEV">
647639
<span class="sig-prename descclassname"><span class="pre">tilelang.env.</span></span><span class="sig-name descname"><span class="pre">DEV</span></span><em class="property"><span class="w"> </span><span class="p"><span class="pre">=</span></span><span class="w"> </span><span class="pre">False</span></em><a class="headerlink" href="#tilelang.env.DEV" title="Link to this definition"></a></dt>
@@ -1052,10 +1044,9 @@ <h3>Benefits<a class="headerlink" href="#benefits" title="Link to this heading">
10521044
<li><a class="reference internal" href="#tilelang.env.COMPOSABLE_KERNEL_NOT_FOUND_MESSAGE"><code class="docutils literal notranslate"><span class="pre">COMPOSABLE_KERNEL_NOT_FOUND_MESSAGE</span></code></a></li>
10531045
<li><a class="reference internal" href="#tilelang.env.TL_TEMPLATE_NOT_FOUND_MESSAGE"><code class="docutils literal notranslate"><span class="pre">TL_TEMPLATE_NOT_FOUND_MESSAGE</span></code></a></li>
10541046
<li><a class="reference internal" href="#tilelang.env.TVM_LIBRARY_NOT_FOUND_MESSAGE"><code class="docutils literal notranslate"><span class="pre">TVM_LIBRARY_NOT_FOUND_MESSAGE</span></code></a></li>
1055-
<li><a class="reference internal" href="#tilelang.env.SITE_PACKAGES"><code class="docutils literal notranslate"><span class="pre">SITE_PACKAGES</span></code></a></li>
1047+
<li><a class="reference internal" href="#tilelang.env.TL_ROOT"><code class="docutils literal notranslate"><span class="pre">TL_ROOT</span></code></a></li>
10561048
<li><a class="reference internal" href="#tilelang.env.TL_LIBS"><code class="docutils literal notranslate"><span class="pre">TL_LIBS</span></code></a></li>
10571049
<li><a class="reference internal" href="#id0"><code class="docutils literal notranslate"><span class="pre">TL_LIBS</span></code></a></li>
1058-
<li><a class="reference internal" href="#tilelang.env.TL_ROOT"><code class="docutils literal notranslate"><span class="pre">TL_ROOT</span></code></a></li>
10591050
<li><a class="reference internal" href="#tilelang.env.DEV"><code class="docutils literal notranslate"><span class="pre">DEV</span></code></a></li>
10601051
<li><a class="reference internal" href="#tilelang.env.THIRD_PARTY_ROOT"><code class="docutils literal notranslate"><span class="pre">THIRD_PARTY_ROOT</span></code></a></li>
10611052
<li><a class="reference internal" href="#id1"><code class="docutils literal notranslate"><span class="pre">DEV</span></code></a></li>

0 commit comments

Comments
 (0)