mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-25 07:14:16 +08:00
6873e19da2b267d3cb7d0305561ba0bed63feae3
4
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Lei Wang
|
5d729eeebc |
[Refactor] Remove GEMM v1 and promote gemm_py to be the canonical gemm op (#2033)
* [Refactor] Remove GEMM v1 and merge gemm_py into the canonical gemm op
GEMM v2 has been the default for a while and no tests or examples still
depend on v1. This removes the legacy C++ GemmNode path entirely and
promotes the former GemmPyNode (Python-delegated lowering) to be the one
and only GEMM tile op.
C++:
- Merge src/op/gemm_py.{h,cc} into src/op/gemm.{h,cc}. GemmPyNode becomes
GemmNode, GemmPy becomes Gemm, FFI type key "tl.GemmPy" becomes "tl.Gemm",
and the op names "tl.tileop.{gemm_py,wgmma_gemm_py,tcgen05_gemm_py}"
collapse to "tl.tileop.{gemm,wgmma_gemm,tcgen05_gemm}". The shared
infrastructure (GemmWarpPolicyNode, GemmInst enum,
GemmWarpPolicyNode::computeWarpPartition) is preserved in place so
GemmSP/GemmSPPy keep compiling unchanged.
- Delete the tl_gemm builtin, its codegen dispatch in
codegen_{cuda,hip,cutedsl}.cc, and its matching branches in
inject_fence_proxy.cc and merge_shared_memory_allocations.cc. tl_gemm_sp
is kept as the remaining async-proxy gemm builtin.
- Drop #include "gemm_py.h" and GemmPyNode/GemmPy references from
inject_pipeline.cc, producer_consumer_ws.cc, and lower_blackwell_2sm.cc.
Python:
- tileop/gemm: rename GemmPy -> Gemm, register as "tl.Gemm", and rename
the FFI callbacks tl.gemm_py.{lower,infer_layout} -> tl.gemm.{lower,
infer_layout}. _ffi_api.GemmPyGemmInst becomes GemmGetGemmInst.
- language/gemm_op.py: delete gemm_v1/gemm_v2 wrappers. gemm() now emits
tl.tileop.gemm directly; wgmma_gemm/tcgen05_gemm drop the _py suffix.
- env.py: remove TILELANG_USE_GEMM_V1 and the use_gemm_v1() helper.
- ir.py: drop the empty v1-era `class Gemm(Node, Scriptable): ...`
placeholder that was colliding with the real tileop.gemm.Gemm wrapper
under the "tl.Gemm" FFI type key.
- contrib/cutedsl: delete cutedsl/gemm_v1.py (runtime helper for the v1
cutedsl lowering path) and its re-export in cutedsl/__init__.py.
Tests / maint:
- testing/python/tilelibrary/test_tilelang_tilelibrary_gemm.py switches
from T.gemm_v2 to T.gemm.
- testing/python/transform/test_tilelang_transform_inject_fence_proxy.py::
test_lower_fence_proxy is rewritten to exercise tl.tl_gemm_sp instead
of the removed tl.tl_gemm, preserving the fence-proxy assertion.
- Delete maint/gemm_v2/ (v1-vs-v2 benchmarking scripts) and drop the
v1-specific comment from maint/scripts/run_local_ci_test.sh.
C++ headers under src/tl_templates/{cuda,hip,cpp}/gemm*.h are left in
place as reference implementations; they are no longer reachable from
the generated code path since the only consumer (tl_gemm builtin) is
gone, but keeping them avoids a large, hard-to-reverse deletion.
* [Fix] Restore cutedsl/gemm_v1.py (Gemm_SM80 / Gemm_SM90 helpers)
The previous commit deleted cutedsl/gemm_v1.py on the assumption that it
was the CuteDSL lowering helper for the removed T.gemm_v1 path. It is
actually a runtime helper module that defines SM80/SM90 CUTE DSL
primitives (make_smem_layout_AB, get_tma_atom, etc.) which are emitted
into generated cutedsl kernel code via
`import tilelang.contrib.cutedsl as tl; tl.Gemm_SM90.get_tma_atom(...)`
(see jit/adapter/cutedsl/wrapper.py CUBIN_TMA_ATOM_INIT_TEMPLATE).
The "v1" in the filename refers to the CuteDSL helper version, not to
the deleted T.gemm_v1 op, so removing it broke the cutedsl code path for
attention_sink, blocksparse_attention, blocksparse_gemm, and deepgemm
examples with `module 'tilelang.contrib.cutedsl' has no attribute
'Gemm_SM90'`. Restore the file and its `from .gemm_v1 import *` wildcard
re-export.
* [Refactor] Restore tl_gemm builtin as a call_intrin path to tl:: templates
The previous commit in this PR deleted the `tl_gemm` builtin entirely on
the argument that no code emitted it anymore. That was too aggressive:
keeping the builtin preserves a deliberate extension point for any
Python-side gemm lowering backend (for example a future
GemmLegacyExtern) to delegate into the existing C++ templates in
src/tl_templates/<target>/gemm*.h via
T.call_intrin("handle", tir.op.Op.get("tl.tl_gemm"),
op_instance_str, A_ptr, B_ptr, C_ptr)
without any extra C++ plumbing. It is also the only thing that keeps
those tl_templates headers reachable rather than dead code.
Restore:
- tl_gemm builtin declaration in src/op/builtin.h (with an updated
docstring describing the call_intrin entry point) and its
TIR_DEFINE_TL_BUILTIN registration in src/op/builtin.cc.
- Codegen dispatch branches in codegen_cuda.cc, codegen_hip.cc, and
codegen_cutedsl.cc that turn a tl_gemm Call into a PrintCallExtern
emission of `tl::gemm_xx<...>(A, B, C)`.
- SMEM alignment and async-proxy fence handling for tl_gemm in
merge_shared_memory_allocations.cc and inject_fence_proxy.cc. These
passes are tied to the builtin's semantics (not to the removed v1
GemmNode) and are still correct the moment anything emits tl_gemm.
- test_lower_fence_proxy in test_tilelang_transform_inject_fence_proxy
reverts to exercising `tl.tl_gemm` instead of the temporary
`tl.tl_gemm_sp` substitution introduced when the builtin was gone.
Note that the current sole gemm lowering path (the merged GemmNode that
delegates to Python via tl.gemm.lower) still emits inline mma/wgmma/
tcgen5mma via the Python macro generators, not tl_gemm. So restoring
the builtin is a no-op for the default path - it only re-enables the
"call into v1 templates from Python" option for anyone who wants it.
* [Fix] Default HIP execution backend to tvm_ffi, not cython
resolve_execution_backend() only promoted cuda and metal to tvm_ffi for
the "auto" path, so hip fell into the else branch and picked cython. But
allowed_backends_for_target() already lists tvm_ffi as an allowed
backend for hip, and the cuda default is tvm_ffi for good reasons:
- The cython SourceWrapper path in src/tilelang/jit/adapter/wrapper.py
has known gaps (for example T.alloc_global is not handled in
create_dispatch_func, which breaks any test that uses an internally
allocated global scratch buffer).
- libgen.py appends user-supplied compile_flags verbatim to the
hipcc/nvcc command. nvcc accepts --use_fast_math, clang++/hipcc does
not, so a CUDA-style flag leaks into HIP compilation.
Neither of these leak through the tvm_ffi backend. Aligning hip's
default with cuda makes the two paths symmetric and unbreaks
test_alloc_global and test_all_attrs_together_lazy on ROCm without
having to mark the tests skip or rewrite the cython wrapper.
cython is still available on hip via execution_backend='cython' for
anyone who explicitly wants it.
* [Test] Gate CUDA-only language tests with requires_cuda
Three language tests hardcode target='cuda' (or rely on CUDA-only
codegen behavior) but lacked a CUDA availability marker, so they failed
on the ROCm CI runner with "Cannot find global function
target.build.tilelang_cuda" once the earlier pre-existing failures got
past --maxfail=3. Add @tilelang.testing.requires_cuda to:
- test_tilelang_if_range
- test_annotate_min_blocks_per_sm_launch_bounds
- test_tilelang_transpose and test_tilelang_transpose_square
Also replace the hand-rolled __main__ block in test_tilelang_transpose
with tilelang.testing.main() so the skip markers are honored when the
file is run directly.
* [Test] Delete tile library gemm test
Remove testing/python/tilelibrary/test_tilelang_tilelibrary_gemm.py.
|
||
|
Lei Wang
|
c750fb8a98 |
[Enhancement] Update examples and tests for improved type handling functionality (#1448)
* [Enhancement] Update examples and tests for improved type handling and functionality
- Enhanced various example scripts to support new data types and improve compatibility with PyTorch.
- Updated tests across multiple modules to ensure correct functionality with the latest changes in type handling.
- Refactored code in examples to streamline operations and improve clarity, particularly in tensor operations and memory management.
- Added comprehensive tests for new features and fixed existing issues related to type conversions and buffer handling.
* [Refactor] Update accumulation data type to float32 across examples
- Changed accumulation data type from "float" to T.float32 in multiple example scripts to ensure consistency and improve numerical stability.
- This update affects various modules including flash attention, GEMM analysis, convolution, and deepseek MLA examples, enhancing type handling across the board.
* [Refactor] Standardize data type usage across benchmark scripts
- Updated data type definitions in benchmark scripts to use T.float16 and T.float32 consistently, enhancing clarity and type handling.
- Adjusted dtype assignments in matmul functions and configuration setups to align with the new standard.
- Improved overall code consistency and maintainability by ensuring uniform data type usage across various modules.
* [Refactor] Standardize data type usage in templates and scripts
- Updated data type definitions in various templates and scripts to use string representations (e.g., "float16", "int32") instead of T.float16 and T.int32 for improved consistency and clarity.
- Enhanced overall code maintainability by ensuring uniform data type usage across multiple modules, including convolution, elementwise operations, and matrix multiplication templates.
- This change aims to streamline type handling and improve compatibility with existing workflows.
* [Refactor] Standardize data type usage in examples and benchmarks
- Updated data type definitions in various example and benchmark scripts to use T.float16 and T.int32 consistently, enhancing clarity and maintainability.
- Adjusted dtype assignments in kernel functions and configuration setups to align with the new standard.
- Improved overall code consistency by ensuring uniform data type usage across multiple modules, including attention mechanisms, matrix multiplication, and GEMM examples.
* [Refactor] Import dtypes from language.v2 module
- Added import statement for dtypes from the language.v2 module to enhance type handling and maintain consistency across the codebase.
- This change aims to streamline data type management and improve overall code clarity.
* fix
* [Refactor] Standardize data type usage across scripts
- Updated data type definitions in various scripts to use string representations (e.g., "float16", "int8") instead of T.float16 and T.int8 for improved consistency and clarity.
- Adjusted dtype assignments in functions and configuration setups to align with the new standard, enhancing overall code maintainability.
- This change affects multiple modules, including benchmark and attention mechanisms, ensuring uniform data type usage throughout the codebase.
* [Refactor] Update data type handling for consistency and clarity
- Changed string representations of data types in the Hint class to use T.float32 and T.int32 for improved consistency.
- Added new data types "int4" and "int16" to the dtypes module, enhancing type support across the codebase.
- Updated function signatures and assertions in the lop3 and mxfp modules to utilize the new data types, ensuring uniformity in type handling.
- This refactor aims to streamline data type management and improve overall code clarity and maintainability.
* [Enhancement] Improve data type handling and error messaging
- Introduced a mapping for canonical data types to their display strings, enhancing clarity in type representation.
- Updated the dtype creation logic to utilize the new mapping, ensuring more intuitive handling of string inputs.
- Refined error messages in the lop3 module to provide clearer feedback on invalid source formats, improving debugging and user experience.
* [Fix] Correct boolean flag in GEMM SP test case
- Updated the boolean flag in the test_gemm_sp_sm90 function to ensure proper functionality in the test case.
- This change enhances the accuracy of the test and aligns it with expected behavior for the GEMM SP implementation.
* [Refactor] Standardize data type usage across scripts
- Updated data type definitions in various scripts to use T.float16 and T.bfloat16 consistently, enhancing clarity and maintainability.
- Adjusted dtype assignments in function signatures and argument parsing to align with the new standard, ensuring uniform data type usage throughout the codebase.
- This change affects multiple modules, including benchmarks and examples, improving overall code consistency and readability.
* [Refactor] Standardize data type usage in various modules
- Updated data type assignments in multiple scripts to utilize T.float32, T.int8, and T.int32 consistently, enhancing clarity and maintainability.
- Adjusted function signatures and parameter types across benchmarks, examples, and tests to align with the new standard, ensuring uniform data type usage throughout the codebase.
- This change improves overall code consistency and readability, impacting modules related to matrix multiplication, GEMM, and tensor operations.
* [Refactor] Update argument parsing for data types in benchmarks
- Changed argument parsing for data types in benchmark_matmul_intrinsic.py and benchmark_matmul_sp.py to use string representations ("float16", "int8", "float") instead of T.float16 and T.float.
- This update enhances consistency in data type handling across benchmark scripts, improving clarity and maintainability.
* [Refactor] Update data type handling in benchmark and example scripts
- Changed data type arguments in benchmark and example scripts to use string representations ("float16") instead of T.float16 for improved consistency.
- Updated function signatures and argument parsing to align with the new standard, enhancing clarity and maintainability across the codebase.
- This change affects multiple modules related to attention mechanisms and tensor operations, ensuring uniform data type usage throughout the examples.
* [Refactor] Fix data type conversion in multiple scripts
- Corrected the usage of the data type conversion method from dtype..as_torch() to dtype.as_torch() across various benchmark and example scripts.
- This change enhances consistency in data type handling and improves code readability, impacting modules related to attention mechanisms and tensor operations.
* [Refactor] Update float8 data type usage across multiple scripts
- Changed instances of T.float8_e4m3 to T.float8_e4m3fn in various benchmark, example, and test scripts to ensure consistency in data type handling.
- This update enhances clarity and maintainability across the codebase, particularly in modules related to matrix multiplication and tensor operations.
* [Refactor] Enhance float8 data type handling in CUDA code generation
- Updated the handling of float8 data types in the CUDA code generation to include additional float8 variants, improving type conversion logic.
- Adjusted conditions to ensure proper type checks for float8 conversions, enhancing clarity and maintainability in the codebase.
- Modified layout inference to streamline float8 type checks, ensuring consistency across the implementation.
- This change impacts modules related to matrix operations and CUDA code generation, improving overall type handling and conversion accuracy.
* [Refactor] Streamline float8 data type handling in CUDA and related modules
- Enhanced float8 data type handling in CUDA code generation by refining type conversion logic and ensuring consistent type checks.
- Updated layout inference for float8 types to improve clarity and maintainability across the implementation.
- This change impacts modules related to matrix operations and CUDA code generation, improving overall type handling and conversion accuracy.
* [Refactor] Remove unnecessary cache disabling in float8 example script
- Eliminated the call to tilelang.disable_cache() in example_group_per_split_token_cast_to_fp8.py to streamline the code.
- This change enhances clarity and maintainability of the example script without affecting its functionality.
* [Refactor] Update data type usage in debug print tests
- Changed the argument for dtype in the test_debug_print_buffer function from a string representation to the corresponding T.bool type.
- This update enhances consistency in data type handling within the test suite, improving clarity and maintainability.
* lint fix
* Update function parameter types from `str` to `T.dtype` for improved type safety in attention sink and related examples
* Refactor `gemv_alloc_reducer` function signature for improved readability by formatting parameters across multiple lines.
|
||
|
Lei Wang
|
29051439db | [Lint] Phaseout Yapf format and embrace ruff format (#1417) | ||
|
Lei Wang
|
c538d8abf5 |
[Language] Support sequence comparisons (#872)
* Update submodule 'tvm' to latest commit 7a71ee34 * lint fix |