Commit Graph
45 Commits
Author SHA1 Message Date
a84c9b1433 [CuTeDSL] Integrate host codegen call sites (#2221)
[CuTeDSL] Integrate host codegen call sites

Drive CuTeDSL launcher generation from lowered host tvm_call_packed call sites so multi-kernel wrappers preserve host IR order and repeated launches can be represented.

Expose separate device, host, and generated module sources through the CuTeDSL adapter and persist them separately in the kernel/autotune caches.

Normalize bare cutedsl targets through TileLang CUDA arch detection so Hopper runners use sm_90a for CuTeDSL compilation instead of the generic sm_90 Target("cuda") result.

Add CuTeDSL host-codegen integration tests covering multi-call launcher ordering, source accessors, cache restore behavior, and bare target architecture normalization.

Co-authored-by: Jayce <jayce.su@enflame-tech.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
2026-05-21 20:28:41 +08:00
Liu YunuoandLeiWang1999 b9a5907a20 [Example] Migrate eligible examples to eager style (#2010)
* [Example] Migrate eligible examples to eager style

* fix wrong shape annotation in gemm_split_k examples

* fix incorrect dtype in gemm_streamk example

* fix

* revert some changes

* Refactor attention functions to use new parameter structures and improve readability. Updated `example_amd_flash_attn_bwd.py` and `example_amd_flash_attn_fwd.py` to define input parameters more clearly. Enhanced shared memory usage in `flashattn_fwd` and `flashattn_sink` functions across various examples, ensuring compatibility with new configurations. Adjusted kernel definitions and added necessary checks for GPU capabilities.

* lint fix

* FIX

---------

Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
2026-05-21 20:27:06 +08:00
b939fa0137 [TIR][IR] Update to use tirx (#2216)
* [Feature] Implement Hoist Loop-Invariant If Pass

- Added a new pass to hoist loop-invariant if statements out of loops, improving optimization opportunities.
- Introduced classes for collecting and checking conditions, as well as for rewriting statements.
- Integrated the new pass into the optimization pipeline and provided a corresponding API for usage.

* Add tilelang semantics guide to programming guides (#1740)

* Add CallNodeChecker to Loop Unswitching Pass and integrate tilelang cache disabling in benchmark script

- Introduced CallNodeChecker class to identify CallNode expressions in loop conditions, enhancing loop-invariant checks.
- Updated IsLoopInvariant function to reject conditions containing CallNodes, preventing potential side effects.
- Added tilelang.disable_cache() in benchmark_mha_sink_fwd.py to optimize performance during benchmarking.

* lint fix

* lint fix

* Enhance Loop Unswitching Pass with Let Binding Support

- Added support for Let-bound variables in the WrittenBufferReadChecker to improve buffer read checks.
- Introduced UsesLoopVarThroughLetBindings function to check if conditions depend on loop variables through Let bindings.
- Updated IsLoopInvariant function to account for Let bindings when determining loop invariance.
- Enhanced HoistableIfFinder to track Let bindings for variables bound to BufferLoad expressions.
- Added debug print statements in the OptimizeForTarget function to visualize the module state before and after loop unswitching.

* fix

* lint fix

* Enhance Loop Unswitching with CallNode Exclusion Check

- Added CallCheckerExcludingIf class to ensure function calls outside of hoisted if statements are identified, preventing potential synchronization issues during loop unswitching.
- Updated loop unswitching logic to incorporate the new call checker, enhancing safety and correctness.
- Integrated debug print statements in OptimizeForTarget to visualize module state before and after loop unswitching.
- Disabled tilelang cache in the benchmark script for improved performance.

* Update TVM subproject to latest commit and remove debug print statements in OptimizeForTarget function

* Add support for disabling loop unswitching optimization

- Introduced a new configuration option `tl.disable_loop_unswitching` to allow users to disable the loop unswitching optimization.
- Updated the Loop Unswitching pass to check this configuration and return the original function if the option is enabled.
- Added relevant documentation in the PassConfigKey enumeration for clarity.

* update tvm

* Refactor: Update TileLang to use tirx instead of tir

- Removed references to `tvm.tir` and replaced them with `tvm.tirx` across various files, including examples and backend operations.
- Updated target configuration in documentation to reflect the new usage of target config dictionaries instead of CLI-style strings.
- Cleaned up `pyproject.toml` and `load_tvm.cmake` by removing obsolete paths.
- Enhanced examples for dequantization and GEMM operations to utilize the new `tirx` constructs.
- Adjusted various backend operations to ensure compatibility with the new `tirx` namespace.

* pass fix

* Update submodule 'tvm' to latest commit f24c7b19

* Apply formatting fixes

* Update CMakeLists.txt to set HIDE_PRIVATE_SYMBOLS for TileLang and modify target_host initialization in lower.py for compatibility with TVM's target handling.

* Refactor CMakeLists.txt and kernel_cache.py to update target names and library dependencies for TVM integration. Adjusted output targets and library names for better compatibility across platforms.

* Update repair-wheel-command in pyproject.toml to exclude libz3.so with wildcard for better compatibility with dependency packages.

* Fix macOS Z3 dylib lookup

* Move macOS Z3 install-name fix to install step

* Update TVM submodule for runtime import fix

* Update ROCm backend for TVM module factory changes

* Fix build-tree Z3 runtime lookup

* Fix macOS build-tree Z3 lookup

* Support legacy target option strings

* Normalize CUDA target in transform test

* Qualify Metal stream FFI lookup

* Fix PyPI Z3 packaging links

* Update TVM for Windows DLL lookup

* Register Windows dependency DLL directories

* Bump tvm-ffi to match tirx

* tvm update

* Apply lint fixes

* Enhance reduction operations with support for nan propagation and vectorization. Introduce new functions for preferred vectorized size and update existing reduction logic to handle packed operations for bfloat16 and float16 types. Add nan-aware min and max operations in CUDA and ROCm backends, and update related tests to validate functionality.

* Refactor ReduceLowerer: Remove unnecessary buffer allocation logic for packed operations in reduce.h

* Normalize PyPI Z3 dylib before linking

* Drop legacy target option string support

* Fix type reference for barrier annotations in CUDA copy operations. Updated `BufferLoadNode` to `tirx::BufferLoadNode` in `GetBarrier` and `LowerCluster` methods to ensure compatibility with recent changes in the TIR API.

---------

Co-authored-by: Yu Cheng <yu.cheng@pku.edu.cn>
Co-authored-by: Yichen Yan <oraluben@outlook.com>
2026-05-20 15:38:13 +08:00
Lei Wang 9c95a42dfb [Refactor] Phaseout legacy util map_torch_type with T.dtype.as_torch (#2075)
* Remove redundant `map_torch_type` in favor of `T.dtype.as_torch()`

`map_torch_type` was a standalone function that converted dtype strings or
TVM DataType objects to `torch.dtype`. This was redundant since
`T.dtype.as_torch()` (and `T.dtype(x).as_torch()`) already provides the
same functionality with better handling (HIP/CUDA awareness, proper dtype
lookup tables, etc.).

Changes:
- Remove `map_torch_type` from `tilelang/utils/tensor.py` and its re-export
- Make `T.dtype()` idempotent for existing dtype objects (pass-through)
- Add backward-compat for `"e4m3fnuz_float8"` TVM internal name in `as_torch()`
- Replace all usages across production code, tests, and examples
- Rebuild Cython extension to reflect the import change

* Refactor output and accumulation data types in GEMM FP8 example

Updated the `out_dtype` and `accum_dtype` variables in the `example_tilelang_gemm_fp8_intrinsic.py` file to use `T.float32` instead of string literals. This change enhances type consistency and aligns with recent updates in type handling.
2026-04-21 18:27:15 +08:00
a2d6e01866 [Bugfix] Fix double buffer versioning when TMA is used without warp specialization (#1962)
* [Bugfix] Fix double buffer versioning when TMA is used without warp specialization

When TMA is enabled but warp specialization is disabled, both
MultiVersionBuffer and InjectSoftwarePipeline independently added
version dimensions to shared data buffers, resulting in 4x buffer
allocation instead of the intended 2x.

Add a `barrier_only` mode to MultiVersionBuffer so that in the no-WS
path it only versions barrier buffers (for mbarrier parity rewriting),
while data buffer versioning is left to InjectSoftwarePipeline.

Also flatten the if/else branches in OptimizeForTarget into a single
linear flow with a three-way branch (WS / TMA-noWS / noTMA), removing
the duplicated LowerOpaqueBlock call.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* [Enhancement] Introduce TMA attribute handling in LowerTileOp and optimization pipeline

- Added a new PrimFunc attribute `tl.has_tma` to indicate whether TMA operations were generated during the lowering process.
- Updated `OptimizeForTarget` to utilize the `tl.has_tma` attribute for determining the appropriate optimization path, enhancing clarity and efficiency in the pipeline.
- Refactored the `module_has_tma` function to check for TMA operations across all functions in the module, streamlining the decision-making process for TMA-related optimizations.

This change improves the maintainability of the code and ensures that TMA operations are correctly accounted for in the optimization pipeline.

* fix

* remove barrier related test for cutedsl.

---------

Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-24 14:21:57 +08:00
Lei WangandZhiwen Mo f8dc61c02c test: reduce CI runtime for slow Python suites (#1932)
* test: trim CI runtime for slow Python suites

* remove useless test

* Update TVM submodule to latest commit and remove manual device assert test from CI. Replace with a call to tilelang.testing.main().

* lint fix

* fix

* revert test for eager jit

---------

Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
2026-03-18 13:08:30 +08:00
Lei WangandClaude Opus 4.6 ded6a99221 [Feature] Add Producer-Consumer Warp Specialization and T.tma_copy() API (#1909)
* [Feature] Add Producer-Consumer Warp Specialization and T.tma_copy() API

This PR introduces Producer-Consumer Warp Specialization for sm90+ TMA
pipelines and a new T.tma_copy() API for explicit mbarrier management.

Key changes:

1. **ProducerConsumerWarpSpecialized pass** (new):
   Splits pipelined TMA loops into producer (TMA loads) and consumer
   (compute) warp groups with back-pressure barriers for buffer reuse.
   Works with num_stages >= 1. Producer warp (128 threads) handles
   arrive_expect_tx + tma_load; consumer warps handle compute + arrive
   on back-pressure barriers.

2. **T.tma_copy() API** (new):
   Fire-and-forget TMA copy with a required `barrier` parameter. Unlike
   T.copy() which emits producer+wait pairs, T.tma_copy() emits only
   arrive_expect_tx + tma_load. User manages synchronization via
   T.mbarrier_wait_parity().

3. **MultiVersionBuffer barrier expansion**:
   Extended to handle `shared.barrier` scope buffers, expanding them for
   pipelining (1D size multiplication instead of prepending a dimension).
   Auto-computes mbarrier parity as `(k // num_stages) % 2`.

4. **Pass ordering** (phase.py):
   LowerSharedBarrier now runs after MultiVersionBuffer in the TMA path
   so barrier buffers retain their `shared.barrier` scope during expansion.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* fix for annotate reg alloc

* Enhance debugging and performance measurement in examples

- Added print statements to output kernel source for both forward and backward sparse MLA examples, aiding in debugging.
- Updated `sparse_mla_fwd_pipelined.py` to run performance regression tests and print average latency, improving performance tracking.
- Disabled auto cache in `example_mha_sink_fwd_bhsd.py` for better control during testing.

* fix

* Enhance debugging and performance tracking in examples

- Added print statements to output kernel source in various example scripts, including `example_gqa_sink_bwd_bhsd.py`, `example_blocksparse_gemm.py`, and `example_dequant_groupedgemm_bf16_mxfp4_hopper.py`, aiding in debugging.
- Updated default `window_size` in `example_gqa_sink_bwd_bhsd.py` for improved configuration.
- Disabled auto cache in `example_gemm_schedule.py` to enhance control during testing.
- Introduced a new test for preserving protected auto-injected wait groups in the async copy optimization process.

* clean

* Remove unnecessary print statements for kernel source in example scripts

- Eliminated print statements from `example_gqa_sink_bwd_bhsd.py`, `example_blocksparse_gemm.py`, and `example_dequant_groupedgemm_bf16_mxfp4_hopper.py` to clean up output and improve readability.
- These changes streamline the examples while maintaining functionality.

* clean

* lint fix

* Refactor T.tma_copy() and barrier synchronization in examples and tests

- Updated `example_warp_specialize_gemm_copy_1_gemm_0.py` to use `T.tma_copy()` for optimized memory copying with explicit barrier management.
- Enhanced test cases in `test_tilelang_issue_tma_no_ws.py` and `test_tilelang_language_tma_copy.py` to reflect changes in barrier synchronization, ensuring proper functionality with the new T.tma_copy() API.
- Removed unnecessary print statements and adjusted barrier allocations for clarity and performance.
- Cleaned up the `test_tilelang_transform_inject_tma_barrier.py` file by deleting it, as it was no longer needed.

* Add buffer data mapping for remapped buffers in multi_version_buffer_rewriter

- Introduced a check to ensure that the data variable of a buffer is discoverable, allowing the barrier_init annotation update to find the remapped buffer correctly.
- This change enhances the functionality of the buffer remapping process, ensuring proper synchronization in the transformation pipeline.

* Implement TMA im2col load synchronization with mbarrier allocation

- Added support for allocating mbarriers to synchronize TMA im2col loads, enhancing the pipeline's barrier management.
- Introduced logic to handle mbarrier creation and integration into the TMA copy process, ensuring proper synchronization and data transfer.
- Updated the handling of the mbarrier in the TMA copy statement to improve performance and maintain functionality within the thread-gated block.

* Implement TMA im2col load synchronization with mbarrier allocation

- Added support for allocating mbarriers to synchronize TMA im2col loads, enhancing the pipeline's barrier management.
- Introduced logic to handle mbarrier initialization and integration within the TMA copy process, ensuring proper synchronization during memory operations.
- Updated the handling of the mbarrier in the im2col operation to improve performance and maintain consistency in multi-threaded environments.

* disable tma lower for test_cutedsl_barrier

* Refactor code formatting in test_tilelang_jit_cutedsl.py

- Improved readability by adjusting the formatting of the `tilelang.compile` function call.
- Enabled the main testing function in the script to run without being commented out.

* Enhance sparse_mla_fwd.py and thread_storage_sync.cc with code cleanup and new functionality

- Added a call to disable cache in the main execution block of `sparse_mla_fwd.py` for improved performance.
- Commented out unused benchmarking code in `test_sparse_mla_fwd` to streamline testing.
- Introduced new utility functions in `thread_storage_sync.cc` to analyze and optimize shared memory access patterns, enhancing synchronization management.
- Improved readability and organization of the code by restructuring and commenting on key sections.

* Add thread storage synchronization utilities and optimizations

- Introduced functions to identify shared-like scopes, buffers, and pointer variables, enhancing memory management.
- Implemented a `BarrierTransparentStmtAnalyzer` to analyze statement transparency concerning shared memory.
- Developed a `ThreadStorageSyncOptimizer` to optimize sequences of statements by removing unnecessary synchronization calls, improving performance in multi-threaded environments.
- Enhanced handling of shared memory access in various call types, ensuring proper synchronization during memory operations.

* remove optimize thread_storage_sync

* Refactor pipeline planning to enhance cp.async group scheduling

- Introduced a last-use index for cp.async groups to improve scheduling accuracy.
- Removed unnecessary tracking of first consumers, simplifying the scheduling logic.
- Updated the handling of cp.async and commit statements to ensure proper group ordering based on last-use.
- Added a new test to verify that cp.async groups are ordered correctly by their last-use index.

* lint fix

* Add CUDA requirements to TMA copy tests and update main execution call

- Added decorators to require CUDA and a minimum compute version for the `test_tma_copy_pipeline_2_stages` and `test_tma_copy_pipeline_3_stages` functions.
- Updated the main execution block to call `tilelang.testing.main()` instead of directly invoking the test functions, improving test execution management.

* Refactor attention sink examples by removing pipelined implementations

- Updated imports in `benchmark_gqa_sink_fwd.py` and `benchmark_mha_sink_fwd.py` to use the non-pipelined versions of the attention sink functions.
- Deleted the pipelined implementations of `example_gqa_sink_fwd_bhsd_wgmma_pipelined.py` and `example_mha_sink_fwd_bhsd_wgmma_pipelined.py`.
- Cleaned up regression and test files to remove references to the deleted pipelined examples, ensuring consistency across the codebase.

* performance fix

* Enhance example_blocksparse_gemm.py and inject_pipeline.cc with debugging and utility improvements

- Added a print statement to output the kernel source in `example_blocksparse_gemm.py` for debugging purposes.
- Updated `inject_pipeline.cc` to include a new utility function for stripping TMA copy write buffer attributes, improving code clarity and functionality.
- Modified `producer_consumer_ws.cc` to utilize the new utility function, enhancing the handling of TMA barriers.
- Expanded the `_compile_tvm_ffi` function in the test file to accept additional keyword arguments, improving flexibility for test configurations.
- Introduced a new test to verify that pure TMA warp specialization correctly emits barriers, ensuring compliance with expected behavior.

* lint fix

* Refactor guard variable declarations in producer_consumer_ws.cc for clarity

- Changed guard variable declarations from non-const to const references to improve performance and readability.
- This minor adjustment enhances the handling of producer guards in the context of TMA barriers.

* lint fix

* fix

* lint fix

* perf fix

* Add print statement for kernel source and update main execution in example_dequant_gemm_w4a8.py

- Introduced a print statement to output the kernel source for debugging purposes.
- Updated the main execution block to call `run_regression_perf()` and print the latency, enhancing performance measurement capabilities.

* fix

* Add debugging output and update execution flow in example_tilelang_block_sparse_attn.py

- Added a print statement to display the kernel source for better debugging.
- Updated the main execution block to disable cache and call `run_regression_perf()`, printing the resulting latency for performance measurement.
- Enhanced the producer-consumer logic in producer_consumer_ws.cc with new utility functions to manage local variable access and improve code clarity.
- Introduced a new test case in test_tilelang_issue_tma_no_ws.py to ensure consumer-only local initializations do not leak into the producer, verifying correct behavior in TMA contexts.

* Revert "Add debugging output and update execution flow in example_tilelang_block_sparse_attn.py"

This reverts commit 4995504098.

* Enhance producer-consumer logic and add tests for TMA behavior

- Updated the logic in producer_consumer_ws.cc to avoid reserving unused pre-loop barriers, improving efficiency.
- Introduced a new test case in test_tilelang_transform_producer_consumer_ws.py to verify that consumer-only local initializations do not leak into the producer.
- Added assertions in test_tilelang_issue_tma_no_ws.py to ensure correct behavior of TMA contexts.

* Add local access tracking and buffer management in producer_consumer_ws.cc

- Introduced LocalAccessSummary and LocalLiveSet structures to track read/write buffers and variable definitions.
- Implemented BufferDataToBufferCollector and LocalAccessCollector classes for efficient buffer data collection and access summary generation.
- Enhanced the logic for handling TMA copy write buffer data extraction, improving the producer-consumer interaction in TMA contexts.

* fix

* fix

* fix

* Refactor producer-consumer logic and enhance kernel execution in tilelang

- Updated the SparseFlashAttn class to improve kernel invocation clarity by separating kernel definition and execution.
- Enhanced producer-consumer workflow in producer_consumer_ws.cc by introducing optional guards for wait statements, ensuring correct behavior in concurrent contexts.
- Added utility functions to manage loop body lets, improving code readability and maintainability.
- Expanded test coverage in test_tilelang_transform_producer_consumer_ws.py to verify the preservation of guarded waits and backpressure logic.

* Enhance test coverage and execution flow in attention sink examples

- Added pytest timeout decorators to tests in test_example_attention_sink.py and test_example_blocksparse_attention.py to prevent long-running tests.
- Updated the main execution block in test_example_attention_sink.py to disable cache and directly call the sliding window test function, improving test execution clarity.

* update

* lint fix

* lint fix

* Refactor and enhance barrier handling in tilelang transformations

- Removed unused VisitStmt_ method from fuse_mbarrier_arrive_expect_tx.cc to streamline code.
- Introduced ProducerSimtCopyDetector in producer_consumer_ws.cc to identify SIMT copy patterns, improving buffer access tracking.
- Enhanced test coverage in test_tilelang_transform_fuse_mbarrier_arrive_expect_tx.py and test_tilelang_transform_producer_consumer_ws.py to validate new functionality and ensure correct barrier behavior.

* lint fix

* lint fix

* Refactor example scripts and clean up utility functions

- Removed commented-out print statements in example_warp_specialize_gemm_copy_1_gemm_0.py and example_warp_specialize_gemm_softpipe_stage2.py for cleaner code.
- Simplified buffer scope checks in utils.h by removing inline functions and directly using scope comparisons in relevant functions.
- Enhanced buffer handling in lower_tile_op.cc and producer_consumer_ws.cc by streamlining condition checks for buffer types.

* Remove unnecessary cache disabling in tilelang tests and examples

- Eliminated calls to tilelang.disable_cache() in example_gqa_decode_varlen_logits.py, test_tilelang_issue_tma_no_ws.py, and test_tilelang_language_tma_copy.py for cleaner code and improved performance.
- Updated test execution flow in test_tilelang_issue_tma_no_ws.py to directly call the main testing function.

* fix

* Add LocalAccessSummary collection for loop bounds in producer_consumer_ws.cc

- Introduced CollectExpr method in LocalAccessCollector to gather access summaries for expressions.
- Updated pre-loop liveness assignment to include variables used in pipeline loop bounds, ensuring accurate classification of scalar setups.

* enhance

* fix

* enhance

* fix

* fix

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-18 13:08:10 +08:00
Gabriel Wuandyuxic 391cf5b6b9 [Enhancement] GEMM V2 on SM90/SM100 CuTeDSL backend (#1855)
* fix(cutedsl): add rsqrt/rsqrtf support in codegen

- Add rsqrt/rsqrtf to fastmath function mapping
- Implement ieee_frsqrt using tl.rsqrt
- Export rsqrt from contrib/cutedsl/math.py

Fixes: norm/test_rms_norm.py

* fix(cutedsl): add atomic_add support and fix shared.dyn access

- Add atomic_add_elem_op and atomic_add_ret_elem_op handling
- Exclude shared.dyn from direct indexing to avoid out-of-bounds access
  (shared.dyn shape is set to (1,) but accessed with arbitrary indices)

Fixes: gemv/test_example_gemv.py

* fix(cutedsl): add builtin::ptx_ldmatrix support

Implement builtin::ptx_ldmatrix by reusing tl.ptx_ldmatrix_x* functions.

* feat(cutedsl): add ptx_mma support with all MMA types

Add complete PTX MMA intrinsic support for CuTeDSL backend:

- Implement ptx_mma codegen in codegen_cutedsl.cc
- Add GetVarPtr_ helper for local buffer pointer extraction
- Create ptx_mma.py with inline PTX assembly wrappers for:
  - FP16: m16n8k16 -> f16/f32
  - BF16: m16n8k16 -> f32
  - INT8/UINT8: m16n8k32 -> i32
  - INT4/UINT4: m16n8k32 -> i32
  - TF32: m16n8k4, m16n8k8 -> f32
  - FP64: m8n8k4 -> f64
  - FP8 e4m3/e5m2: m16n8k32 -> f16/f32
- Fix BroadcastNode to emit typed values for make_filled_tensor
- Update ptx_ldmatrix to use GetVarPtr_ for local buffers

* feat(cutedsl): comprehensive atomic operations support

- Add atomic.py module with full atomic operations:
  - AtomicAdd: float32 (nvvm.atomicrmw), float16 (inline PTX with i16 bitcast),
    int32/int64 (nvvm.atomicrmw)
  - AtomicAddx2/x4: vectorized float32 via PTX atom.global.add.v2/v4.f32
  - AtomicMax/Min: int32/int64 (nvvm.atomicrmw), float32 (CAS loop via inline PTX)
  - AtomicLoad/Store: all memory orderings via inline PTX (relaxed/acquire/release/seq_cst)
- Add codegen support for atomic_load/store/max/min/addx2/addx4 in codegen_cutedsl.cc
- Move AtomicAdd from __init__.py to atomic.py, extend with float16/int64 support

* feat(cutedsl): add fabsf/fabs support in codegen and math module

* feat(cutedsl): add wait_wgmma support via nvvm.wgmma_wait_group_sync_aligned

* feat(cutedsl): support i32 min/max

* fix(cutedsl): handle FP8 narrow precision alignment in codegen

CuTeDSL requires all FP8 vector operations to be 32-bit aligned, but
TIR generates float8x2 (16-bit) vectors. Additionally, MLIR vector
extractelement on FP8 types fails with unrealized_conversion_cast.

Fix CastNode, BufferLoadNode, and BufferStoreNode to:
- Pad unaligned FP8 casts to aligned width (e.g. f32x4 -> fp8x4)
- Return rmem tensor names (not .load()) so downstream uses rmem
  element access instead of MLIR vector extractelement
- Use padded aligned stores for local rmem, scalar element stores
  for global/shared memory (to handle non-aligned offsets)

* fix(cutedsl): use pointer arithmetic for shared.barrier buffer access

_Pointer type from tl.alloc_smem() does not support subscript [i],
but supports pointer arithmetic (ptr + i). This fixes warp_specialize
tests that use mbarrier.

- GetBufferRef_: use (vid + index) for shared.barrier instead of vid[index]
- BufferLoadNode: skip .load() suffix for shared.barrier (returns pointer, not value)

Fixes 3/4 warp_specialize tests.

* fix(cutedsl): handle single coordinate in tma_load

When tma_load receives a single Int32 coordinate (for 1D tensors) instead
of a tuple, wrap it in a tuple before iterating. This fixes the error:
'Int32' object is not iterable

Fixes flash_decoding/test_example_flash_decoding.py

* feat(cutedsl): add decode_i4u/s_to_f16 and warp primitives

- Add quantize.py with decode_i4u_to_f16 and decode_i4s_to_f16
  using inline PTX (lop3.b32 + sub.f16x2)
- Add warp.py with __activemask, __shfl_down_sync, __shfl_sync
  using CuTeDSL's shuffle_sync functions
- Update codegen to map decode_* and __* functions to tl.* namespace

Fixes dequant_gemv_fp16xint4 test.

* feat(cutedsl): add decode_fp4_to_bf16_twiddling

Implement FP4 to BF16 conversion using bit twiddling technique.
- Use prmt.b32 for endianness handling
- Use mul.bf16x2 for scaling
- Pack output with correct high/low BF16 layout

Fixes test_example_dequant_gemm_bf16_mxfp4_hopper.

* fix(cutedsl): support reinterpret for arbitrary expressions and fix shift type promotion

- reinterpret: handle non-BufferLoad args with tl.bitcast()
- shift_right/shift_left: wrap result in dtype() to preserve type
- CuTeDSL Int8>>4 returns Int32 (Python promotion), CUDA keeps int8_t
- Add bitcast() function using llvm.bitcast MLIR op

* feat: disable v2->v1 fallback

* fix(cutedsl): TensorSSA for store (Select/scalar) and TF32 MMA dispatch

- ArithValue: emit tl.where() for Select and wrap scalar store RHS in
  make_filled_tensor().load() so CuTeDSL Tensor.store() gets TensorSSA.
- TF32: accept fp32/f32/float32 for m16n8k4 and m16n8k8 in ptx_mma dispatch
  (e.g. deepseek_mhc).

* feat(cutedsl): implement CumSum1D and CumSum2D for cumulative sum operations

- Add CumSum1D and CumSum2D classes in reduce.py based on CUDA templates
- Add __shfl_up_sync wrapper in warp.py using CuTeDSL's shuffle_sync_up
- Implement warp-level prefix sum (forward/reverse) using shuffle operations
- Use @cute.jit decorator and nested if for DSL-compatible control flow
- Fixes gdn/test_example_cumsum_compilation test (7/8 gdn tests now pass)

* feat(cutedsl): implement tma_reduce for TMA atomic add reduction

* fix(cutedsl): handle if_then_else as tl.where and unify Select/where arm dtypes

- Override builtin::if_then_else in CallNode to emit tl.where() instead of
  Python ternary (which yields ArithValue, not TensorSSA)
- Use make_rmem_tensor + element assignment instead of make_filled_tensor
  (which creates unsupported vector<1xT> types for int)
- Unify arm dtypes in both SelectNode and if_then_else: when true/false arms
  differ (e.g. BFloat16 vs Float32), cast both to common type via .to()
- Fix BufferStore: skip make_filled_tensor wrapping for Select/if_then_else
  values (already TensorSSA); extract [0] for scalar element assignment
- Fix GetBufferRef_: extract Ramp base for make_tensor_at_offset offset
  parameter to avoid emitting tuple offsets

Fixes: cast/group_per_split_token, topk, deepseek_mhc, flash varlen

* fix(cutedsl): convert non-integer offsets to Int64 in make_tensor_at_offset

Complex arithmetic on cutlass.Int64 values (involving tl.min, //, %) can
produce Float64 ArithValue that CuTeDSL Pointer.__add__ does not accept.
Detect non-integer offsets via is_integer() and wrap with cutlass.Int64().

Fixes blocksparse_attention varlen_indice and varlen_mask (6/6 passed).

* fix(cutedsl): fix narrow unaligned load/store for sub-32-bit types

Three issues in the narrow unaligned (e.g. uint8x2, fp8x2) code paths:

1. Load from shared/global: using aligned_lanes (4) elements caused OOB
   near buffer boundaries and MisalignedAddress for non-aligned offsets.
   Fix: load exactly value_lanes (2) elements with no cute.assume.

2. Load from rmem: cute.assume(offset, divby=aligned_lanes) silently
   truncates offsets not divisible by aligned_lanes (e.g. offset=6 with
   divby=4 → 4). Fix: drop cute.assume for rmem (no hardware constraint).

3. Store to rmem: padded (aligned_lanes,) .store() writes beyond
   value_lanes elements, causing overlapping writes and OOB at tensor
   boundaries. Also triggers CuTeDSL's 32-bit alignment check for FP8.
   Fix: use element-by-element assignment (vid[i]=val).

Fixes dequantize_gemm w4a8 (cudaErrorMisalignedAddress → correct results).

* fix(cutedsl): handle signless int types in tl.min/tl.max fallback

* fix(cutedsl): handle _Pointer src_values in AtomicAddx2/x4

* fix(cutedsl): remove double PrintIndent in cp_async_gs_conditional codegen

* fix(example): support SM100 in attention_sink bwd configs

Change get_bwd_configs() to use >= 90 instead of == 90 so that SM100
(Blackwell) uses the same config as SM90 (Hopper): block_M=128,
block_N=32, num_stages=2, threads=256.

* feat: implement WGMMA V2 (ss + rs) for CuTeDSL backend

Implement explicit WGMMA descriptor management ops in the CuTeDSL backend,
replacing LOG(FATAL) stubs with working codegen for all 8 WGMMA-related ops.

New Python files:
- typing.py: WGMMA type map (string dtype -> nvvm.WGMMATypes)
- gemm_V2.py: GmmaDescriptor class, initialize/increase descriptor,
  warpgroup sync ops, wgmma_ss (via nvvm.wgmma_mma_async),
  wgmma_rs (via PTX inline asm)

Key findings and fixes:
- nvvm.mma_smem_desc does not correctly pack the layout_type/swizzle
  field in the WGMMA descriptor; replaced with manual bit packing
- wgmma_rs implemented using llvm.inline_asm with dynamically generated
  PTX and constraint strings based on M/N/K Constexpr parameters

C++ codegen changes (codegen_cutedsl.cc):
- warpgroup_arrive/commit_batch/wait/fence_operand
- ptx_wgmma_ss (15 args) and ptx_wgmma_rs (14 args)
- initialize_wgmma_descriptor (arg reorder: TIR->Python)
- increase_descriptor_offset

All GEMM tests (4/4) and flash attention tests (14/14) pass.

* fix(cutedsl): replace break with guard variable pattern for CuTeDSL compatibility

* feat(cutedsl): implement tcgen05 (SM100/Blackwell) MMA support

Add CuTeDSL codegen and Python runtime support for tcgen05 MMA operations:

- New gemm_tcgen05.py: Tcgen05SmemDescriptor, descriptor initialization,
  MMA variants (ss/ws_ss/ts), mbarrier arrive, TMEM alloc/dealloc/load
- All MMA and commit instructions guarded by elect_one_sync via inline PTX
  (@q predicate on individual instructions, not on block scopes)
- TMEM load uses recursive power-of-2 splitting with direct MLIR emission
- C++ codegen: replace 6 LOG(FATAL) stubs with proper code generation for
  tcgen05 MMA/descriptor/TMEM ops, handle descriptor scope allocation
- Fix shared.barrier pointer access (skip .iterator for alloc_smem results)
- Fix mbarrier_wait: use blocking PTX loop instead of single try_wait call

* fix(cutedsl): correct tcgen05_ld 64b/128b/256b PTX variants and cleanup

- Fix tcgen05_ld_32dp64bNx/128bNx/256bNx: use correct PTX instructions
  (16x64b/16x128b/16x256b) with two-half decomposition instead of
  hardcoded 32x32b
- Generalize _emit_tmem_ld_segment to accept ptx_type and regs_per_x
- Remove dead expression in ptx_mma.py (c_ptr + a_offset)
- Remove unused GemmCuTeDSL and is_cutedsl_target imports

* fix(cutedsl): review cleanup — warp shuffle, mbar timeout, atomic addr space, dead code

- Fix warp.py: correct mask_and_clamp formulas to match CUDA semantics
  for __shfl_down_sync, __shfl_up_sync, __shfl_sync (width param was ignored)
- Fix mbar.py: use timeout_ns parameter instead of hardcoded 0x989680
- Fix atomic.py: use generic address space (ld.f32, atom.cas.b32) instead of
  .global for AtomicMax/AtomicMin CAS loops, supporting shared memory pointers
- Add ICHECK for non-unit Ramp stride in codegen_cutedsl.cc GetBufferRef_
- Remove unreachable duplicate elif gemm_inst.is_tcgen5mma() branch
- Remove unused _to_ir_value helper from ptx_mma.py

* fix(cutedsl): skip BulkLoad1D for CuTeDSL + fix uint broadcast type mismatch

Two fixes for dequant/groupedgemm_bf16_mxfp4_hopper on CuTeDSL:

1. Skip BulkLoad1D/BulkStore1D for CuTeDSL target in copy.cc:
   cp_async_bulk_shared_cluster_global (raw 1D TMA) combined with WGMMA
   in the same kernel triggers a ptxas ICE in the NVPTX backend.
   Falls through to descriptor-based BulkLoad/BulkStore instead.

2. Use signed int type for BroadcastNode fill values in codegen_cutedsl.cc:
   CuTeDSL/MLIR normalizes unsigned integer tensor loads to signed types
   (Uint8 pointer -> i8 elements). BroadcastNode fill values now use
   matching signed types to avoid Int8/Uint8 mismatch in tl.where().

* fix(cutedsl): support fp16/bf16 in AtomicAddx2 via atom.add.noftz.v2.f16/bf16

* ci(cutedsl): add CuTeDSL examples to CI pipeline

Add a new GitHub Actions step that runs all examples with
TILELANG_TARGET=cutedsl. Known failures (sparse ops) are marked as
xfail and unsupported tests (convolution, stream-K, flash_decoding FA3)
are auto-skipped via pytest_collection_modifyitems in conftest.py.
The step uses continue-on-error: true initially until stabilized.

* chore: fix pre-commit

* fix(adapter): use name-based fallback for tir.Var dynamic symbolic lookup

tir.Var uses object identity for hash/eq, so when param_shapes (from
the original PrimFunc) and dynamic_symbolic_map (from the transformed
PrimFunc in ir_module) contain different tir.Var objects with the same
name, direct dict lookup fails with KeyError.

Add a secondary name-based index (_dynamic_symbolic_name_map) and a
_lookup_dynamic_symbolic() helper that falls back to name matching.
This fixes KeyError: batch under pytest-xdist parallel execution.

Affected adapters: cutedsl, nvrtc. (cython and tvm_ffi already had
a str()-based workaround.)

* chore: remove dead GemmCuTeDSL class and update stale V1 comments

- Delete `gemm_cutedsl.py` (no references anywhere in codebase)
- Remove outdated comment in `gemm_op.py`
- Update `run_local_ci_test.sh` comment: CuTeDSL no longer requires V1
- Keep `TILELANG_USE_GEMM_V1` env var as a debugging escape hatch

* refactor(cutedsl): reorganize contrib/cutedsl module structure

- Rename gemm_V1.py/gemm_V2.py to snake_case (gemm_v1.py/gemm_v2.py)
- Merge typing.py (19 lines) and mbar.py (53 lines) into utils.py and cpasync.py
- Extract utility functions from __init__.py into new utils.py module
- Slim __init__.py to pure re-exports
- Add __all__ to all submodules for explicit public API documentation

* fix(ci): use .empty() instead of .size() >= 1 in pipeline_planning.cc

* fix: coderabbit comments

* fix(cutedsl): correct WGMMA result register count for f16/bf16 output

wgmma_ss/wgmma_rs used M*N//128 for result registers, which is only
correct for 32-bit output types. For f16/bf16, each i32 register holds
2 packed elements, halving the count. Derive from cutlass type .width
instead of hardcoding.

Also update barrier test assertion to match codegen's pointer-arithmetic
syntax: (barriers + i) instead of barriers[i].

* refactor(cutedsl): templatize ptx_mma.py with factory function

Replace 13 near-identical @dsl_user_op PTX MMA functions (~840 lines)
with a single _make_ptx_mma factory and a config table (~120 lines).
Also correctly interpolates a_layout/b_layout into the PTX mnemonic
instead of hardcoding row.col.

* fix(cutedsl): fix NaN CAS loop, wgmma_rs int accum, and stale pipeline state

- atomic.py: Use integer-domain comparison (setp.ne.b32) in AtomicMax/Min
  CAS loops to prevent infinite spin when values are NaN
- gemm_v2.py: Use "=r"/i32 constraints for wgmma_rs when C_dtype is s32
  (int8 GEMM) instead of hardcoded "=f"/f32
- reduce.py: Add assert H <= 32 guard for CumSum2D dim=0 column-wise path
- pipeline_planning.cc: Clear pending tcgen05 state unconditionally after
  arrive to prevent stale entries from leaking to future arrives

* fix(cutedsl): fix ptx_mma brace escaping and wgmma_rs DSL compatibility

- ptx_mma.py: Double-brace register lists in PTX template so
  .format(a_layout=, b_layout=) doesn't interpret {$0, ...} as
  format fields (KeyError regression from templatization)
- gemm_v2.py: Replace if/else branches inside DSL-traced loops with
  ternary expressions to avoid CuTeDSL "operands type change inside
  dynamic if" error
- ci.yml: Remove continue-on-error from cutedsl CI, reduce maxfail
  to 3, increase parallelism to 4

* feat(cutedsl): implement ptx_mma_sp (sparse MMA) codegen and runtime

* fix(cutedsl): guard saturate param, move CuTeDSL CI step after core tests

- ptx_mma/ptx_mma_sp: raise NotImplementedError if saturate=True
- warpgroup_fence_operand: add comment explaining no-op rationale
- ci.yml: move CuTeDSL examples after CUDA/ROCm/Metal test steps
  so a CuTeDSL failure no longer blocks core tests

* feat(cutedsl): implement 16 previously unsupported ops

Add support for thread index queries (get_lane_idx, get_warp_idx,
get_warp_idx_sync, get_warp_group_idx), exp10 math function,
warp-level reductions (sum, max, min, bitand, bitor), and
IEEE-754 compliant arithmetic with explicit rounding modes
(add, sub, mul, fma, rcp, sqrt, div).

Reduces unsupported op count from 26 to 10. The remaining ops
are either high-complexity (tma_load_im2col, tl_gemm_sp),
require wrapper changes (sync_grid), or are legacy/unused.

* fix(cutedsl): wrap non-int32 integer literals with CuTeDSL type constructors

CuTeDSL's as_numeric() converts bare Python ints to Int32, causing
"Type mismatch" errors when storing to narrower tensors (e.g. Int16).
Override IntImmNode to emit cutlass.Int16(0) instead of bare 0 for
non-int32 literals. Also add a BufferStore narrowing guard as safety net.

Fixes compress_kernel Int32→Int16 error in gemm_sp/custom_compress.

* feat(cutedsl): implement sync_grid with cooperative kernel launch

Add grid-level synchronization via a software barrier using
llvm.mlir.global counter + inline PTX atomic spin-wait. The CuTeDSL
wrapper detects sync_grid usage and automatically switches to
cuLaunchCooperativeKernel to guarantee all blocks are resident.

- grid_sync.py: module-global counter + PTX barrier protocol
- wrapper.py: CPP_COOPERATIVE_KERNEL_LAUNCH_TEMPLATE
- codegen_cutedsl.cc: emit tl.sync_grid() instead of LOG(FATAL)

---------

Co-authored-by: yuxic <yuxic@nvidia.com>
2026-02-23 01:22:47 +08:00
fe5bf3569d [Refactor] Enhance T.alloc_barrier with new features and deprecate legacy mbarrier related intrinsics (#1733)
* enhance `T.alloc_barrier` and let `T.wait/arrive` only accept buffers

* lint

* Refactor mbarrier handling and update related documentation

- Updated the `T.alloc_barrier` function to accept an array of arrive counts instead of individual thread counts.
- Modified the `get_mbarrier` function documentation to clarify its purpose and usage.
- Enhanced the `CreateMBarrierBuffer` function to create mbarrier buffers with shared storage scope.
- Updated tests to reflect changes in mbarrier allocation and usage.
- Adjusted comments in various files to improve clarity regarding mbarrier functionality.

* lint

* fix test

* tcgen5 check input

* lint

* fix name conflict

* add name check

* apply ai suggestions

* fix shared mem disalignment and missing fence for mbarrier init

* fix cutedsl codegen

* bugfix

* add test for cutedsl codegen

* Update testing/python/jit/test_tilelang_jit_cutedsl.py

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

* Remove tilelang.disable_cache() call

Remove cache disabling from the test file.

---------

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>
2026-01-28 15:20:13 +08:00
Gongen-Ali 4e53b50750 [Bugfix] Open Rocm ci test and fix some bugs. (#1443)
* Fix some bugs on ci and open rocm ci test

* fix pre-commit

* fix build issue on hip

* refactor on hip

* pre-commit fix

* fix test with pdl

* fix pdl test
2026-01-06 02:18:29 +08:00
9fd6ea7d3b [Feat] PDL Support (#1494)
* [Feat] PDL Support

* remove unused comments/codes

* fix mismatch funcname

* Resolve duplicate config issue for AI review

* fix typo error && clang-tidy error

* chore: retrigger CI

* change file comment

* rename test_tilelang_jit_ctypes.py to test_tilelang_jit_cython.py

* remove ctype backend

* Move pdl attributes to header and remove redundant VisitStmt_

* chore: trigger CI

* modify the cuda kernel with pdl due to nvcc's bug

* throw an error when invoking __ldg with pdl due to nvcc's bug

* Throw an error when pdl is not supported

* fix checking about pdl

* add comments about MarkCudaSyncCalls

* remove pdl support with compute_capability<90

* ruff check

* fix nvrtc about pdl

* update tvm & extend pdl support for tvm_ffi & add test for tvm_ffi

* ensure that kUseDynamicSharedMemoryTag is the last tag in launch_param_tags

* fix bug introduced by branch merging

* [Feature] Add get_host_source method to CythonKernelAdapter

* Introduced a new method, get_host_source, in the CythonKernelAdapter class to return the source code of the host function. This addition enhances the functionality of the adapter by providing direct access to the host kernel source.

---------

Co-authored-by: silentCoder-dev <silentcoder@foxmail.com>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
2026-01-04 17:03:37 +08:00
Gabriel Wu 72ce848be0 [Feat] Adapt gemm v2 for cutedsl backend (#1544)
* feat: adapt gemm v2 for cutedsl backend

* fix: ruff
2025-12-27 13:07:16 +08:00
Jinjie LiuandZihua Wu 9ff7c523a9 [Enhancement] Refactor KernelCache to use inheritance-based design (#1483)
* refactor kernel_cache

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* set key to auto when execution_backend is None

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* fix type error on _load_kernel_source override

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* put kernel cache subclasses in adapter

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* remove auto in execution_backend

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* remove _save_kernel_source_code_to_disk for NVRTCKernelCache

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* refactor all cutedsl if statements

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* remove unused kernel_cubin_path

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* add func arguments

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* remove ctypes

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* update ci.yaml to test kernel cache

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* fix bugs after #1512 merged

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

* test: fix kernel cache test

* fix: ruff

* fix: ruff

* fix: coderabbit

* fix: coderabbit

* fix: ruff

* fix: MMA

* fix: coderabbit

* fix: codex

* fix: coderabbit

* fix: coderabbit

---------

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>
Co-authored-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
2025-12-26 15:41:22 +08:00
Lei Wang 09385e7deb [Refactor] Support auto swizzling for tma store and phaseout related layout annotations (#1509)
* Remove unnecessary swizzled layout annotations from various attention sink examples and kernel analysis script for improved clarity and performance.

* Remove unnecessary swizzled layout annotations from various examples to enhance code clarity and performance.

* lint fix

* Uncomment main test execution for vectorized cast

* Add swizzled layout annotation for B_shared in dequantize GEMM example

This change introduces a layout annotation for the B_shared tensor in the example_dequant_gemm_bf16_mxfp4_hopper.py file, enhancing the memory layout optimization for better performance during matrix multiplication.

* Remove unnecessary cache disabling in GQA sink example for improved clarity and performance.

* lint fix

* Refactor layout functions to support row-major linear layout for any dimension

- Renamed `makeGemmLayoutLinear` to `makeLinearLayout` and updated its implementation to handle arbitrary dimensions.
- Updated related function calls in `gemm_layouts.cc` and `layout.cc` to use the new layout function.
- Enhanced layout inference in `CumSumOpNode` to enforce linear layout for shared buffers in strict mode.

* Refactor `make_linear_layout` to accept a single argument and support arbitrary dimensions

- Updated the function signature to take a `Buffer`, `BufferLoad`, or `BufferRegion` directly.
- Simplified the implementation by removing argument checks and directly obtaining the shape from the buffer info.
- Enhanced the documentation to clarify the function's purpose and usage.

* skip callback test

* Remove example_dequant_gemm_bf16_mxfp4_hopper_tma.py file, eliminating unused code related to dequantization GEMM example.

* typo fix
2025-12-24 14:22:35 +08:00
Kuris c7e8cab883 [Testing] Add Memory Leak Test (#1516)
* fix TL_ENABLE_PTXAS_VERBOSE_OUTPUT is not print in tvm-ffi

* add memory leak test in tvm-ffi

* fix lint error

* fix typos
2025-12-24 14:19:39 +08:00
Lei Wang 11f122e6a8 [Refactor] Phaseout execution_backend ctypes (#1510)
* Refactor execution backend options by removing 'ctypes' from the list of supported backends across multiple files. Update related documentation and tests to reflect this change, ensuring consistency in the autotuning and JIT compilation processes.

* minor fix
2025-12-23 19:52:16 +08:00
7248a810d9 feat(cutedsl): add CuTeDSL backend (#1421)
* feat: CuTeDSL backend

* fix: clang-tidy

* fix: clang-format

* fix: ci

* fix: revert example gemm fp8

* fix: remove duplicate code

* fix: switch-case

* fix: fp16 silence

* fix: TVM IR print

* fix: useless tir

* fix: clang-format

* fix: remove tilelang/contrib/cutedsl/.gitignore

* fix: use hexfloat

* fix: gsym guard

* fix: unknown storage sync type

* fix: string literal

* fix: add args guard

* fix: name hint dedup

* fix: better find_kernel_by_pattern

* fix: set libpath for from_database path

* fix: guard buffer.strides

* fix: from guard

* fix: eviction guard

* fix: use thread local tma descs

* fix: ruff

* fix: drop tma_init_cpp

* fix: exc_info

* fix: negative unmatch early return

* fix: rename postproc func and add test

* fix: handle fast math according to pass config

* fix: dyn_sym parse

* fix: wrap_forward

* fix: use tvm_ffi.libinfo instead of cli

* fix: keep signature

* fix: C++ string safety

* fix: mark tma_store_add as unsupported

* fix: tvm version

* resolve ldsm and cpasync issues.

* fix: minor fixes

* fix: parse signature using ast

* fix: guard global_addr

* fix: create tempfile only when necessary

* fix: use logger.execption for exceptions

* fix: guard lib_path and host_func

* fix: remove tma_cpp_init and add timeout for cpp compile

* add timeout for mbarrier_wait.

* fix: _load_kernel_from_disk signature

* resolve codegen issues.

* fix: logger.exception

* add comment for div_by=1

* merge

* fix: reserve cutlass,cute,tl

* fix: guard tma_store

* fix: allow int64 offset in make_tensor_at_offset

* fix: guard barrier

* fix: add comments for div_by=16

* fix: div_by=1 issue

* delete div_by when offset is 0

* use tl.make_tensor when offset is 0

* fix: explicitly check cutedsl target

* fix: use param.torch_dtype()

---------

Co-authored-by: yuxic <yuxic@nvidia.com>
Co-authored-by: Yong <yong@local>
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
2025-12-18 11:55:18 +08:00
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.
2025-12-17 14:41:25 +08:00
Lei Wang 29051439db [Lint] Phaseout Yapf format and embrace ruff format (#1417) 2025-12-12 17:22:18 +08:00
Lei Wang 1e92d11cd2 [Refactor] Improve assertion handling in CodeGenCHost and ArgBinder (#1352)
* [Refactor] Improve assertion handling in CodeGenCHost and ArgBinder

This commit refines the assertion message generation in CodeGenCHost by optimizing the handling of equality checks and reducing buffer size for error messages. Additionally, it enhances the ArgBinder by introducing a nullable guard mechanism for assertions, allowing for more precise error handling when binding arguments. The changes improve the clarity and efficiency of assertion handling across the codebase.

* [Enhancement] Update matmul kernel and optimize argument binding

This commit enhances the matmul kernel by introducing additional tensor parameters and refining the pipeline stages for improved performance. It also updates the argument binding mechanism to include a flag indicating whether buffers are used, enhancing the efficiency of buffer management. Furthermore, the optimization phase in the engine is improved by adding a simplification step, ensuring better performance and clarity in the generated code.

* lint fix

* [Enhancement] Add tensor checks documentation and improve argument binding assertions

This commit introduces a new documentation page for host-side tensor checks, detailing the automatic validations performed by TileLang on kernel arguments. It enhances the ArgBinder by adding assertions for non-null pointers when arguments are used, improving error handling. Additionally, the optimization phase in the engine is updated to include a simplification step, ensuring better performance and clarity in the generated code.

* [Enhancement] Update .gitignore and refine matmul kernel for improved performance

This commit adds host checks logs to the .gitignore file to prevent unnecessary log files from being tracked. Additionally, it refines the matmul kernel by adjusting pipeline stages, updating tensor parameters, and enhancing argument handling for better performance. The changes also include improved error messages in the argument binding process, ensuring clearer diagnostics for users.

* lint fix

* lint fix

* [Refactor] Simplify tensor_null_test function and remove ptr_null_test

This commit refactors the tensor_null_test function by adding a with_bias parameter and removing the ptr_null_test function, which was previously unused. The run_test function is updated to reflect these changes, streamlining the testing process for tensor operations.

* lint fix

* fix
2025-11-28 03:28:14 +08:00
KurisandLei Wang bccb6485e4 [Feat] Add support for using T.Tensor(n * 2 + 1) in function annotation (#1285)
* [Feature] Add support for A: T.Tensor(n + 1) and A: T.Tensor(2*n)

* issue fix

* fix

* fix

* decreate nproc for debugging

---------

Co-authored-by: Lei Wang <leiwang1999@outlook.com>
2025-11-20 13:56:09 +08:00
Lei Wang 74da369695 [FFI] Use tvm ffi as the default execution backend (#1259)
* [Refactor] Update FFI type handling and simplify argument management

* Refactored FFI type definitions in runtime and code generation files to use `TVMFFIAny` instead of `TVMValue`, enhancing type clarity.
* Updated function registration in `runtime.cc` to utilize canonical names for better consistency.
* Simplified argument handling in the `simplify` transformation, ensuring unused buffer parameters are removed only when simplification is enabled.
* Adjusted autotuner and profiler parameters to standardize the execution backend to `tvm_ffi`, improving clarity in backend selection.
* Removed obsolete `adapt_torch2tvm` function from tensor utilities to streamline the codebase and reduce complexity.

* [Update] Sync TVM submodule and enhance kernel source handling

* Updated the TVM submodule to commit cdc2aced, ensuring compatibility with recent changes.
* Added functionality to print kernel source in `example_blocksparse_gemm.py` for better debugging.
* Commented out the main execution call in test files to prevent unintended execution during testing.
* Introduced `tilelang.disable_cache()` in various test files to streamline testing and avoid cache-related issues.
* Refactored kernel source retrieval methods to improve clarity and consistency across different execution backends.

* [Refactor] Clean up imports and improve code formatting

* Removed unused import of `tilelang.testing` in `test_example_blocksparse_gemm.py` to streamline the code.
* Reformatted several lines in `arg_binder.cc`, `make_packed_api.cc`, `tvm_ffi.py`, and `adapter.py` for improved readability and consistency.
* Updated comments and spacing in `tvm_ffi.py` to enhance clarity without altering functionality.

* Update execution backend options and improve resolution logic

- Changed default execution backend from "cython" to "auto" in multiple locations to allow automatic selection based on the target.
- Expanded the list of supported execution backends to include "torch" and "nvrtc" across various classes and functions.
- Enhanced backend resolution logic in `KernelCache` and `AutoTuner` to ensure appropriate backend selection based on the target.
- Updated documentation to reflect changes in execution backend options and their defaults.

* lint fix

* fix

* Enhance argument handling in CUDA and HIP runtime modules

- Updated `ExtractFuncInfo` in `rt_mod_cuda.cc` and `rt_mod_hip.cc` to map boolean argument types to int32, ensuring compatibility with device runtime.
- Refactored `BindDLTensor` in `arg_binder.cc` to improve null handling and validation checks for DLTensor parameters, utilizing expression-level guards to prevent dereferencing null pointers.
- Enhanced error checking for buffer shape, strides, and data fields, ensuring robust handling of optional inputs and maintaining consistency across various checks.

* lint fix

* lint fix

* lint fix

* lint fix

* minor fix

* fix

* recover check

* Refactor argument binding and validation in `arg_binder.cc`

- Improved null handling and validation checks in `BindDLTensor`, ensuring safe dereferencing of pointers.
- Enhanced consistency checks for buffer shape, strides, and data fields, utilizing expression-level guards.
- Updated `MakePackedAPI` to maintain code clarity and consistency in argument handling.
- Minor adjustments in test files to streamline kernel execution and improve readability.

* lint fix

* stride fix

* minor fix

* fix

* lint fix

* lint fix

* Add CUDA stream access policy window helpers and integrate with L2 persistent cache management

- Introduced functions to set and reset the CUDA stream access policy window, allowing for better control over L2 cache usage.
- Updated runtime files to include new FFI packed functions for managing stream attributes.
- Modified lower_hopper_intrin to incorporate prologue and epilogue statements for L2 cache setup and teardown.
- Enhanced tests to verify the inclusion of new FFI calls in the generated kernel source.

* check with symbolic

* support null ptr

* Update CMakeLists and lower.py for code generation and subproject status

- Added `codegen_c_host.cc` to the list of source files in CMakeLists.txt for improved code generation support.
- Updated the function call in `lower.py` to use `target.build.tilelang_c` for C target host code generation, enhancing compatibility.
- Marked the TVM subproject as dirty to indicate local modifications.

* lint fix

* Update comments for clarity in quickstart.py
2025-11-18 23:50:57 +08:00
Gabriel Wu eb41574431 [fix] NVRTC execution backend (#1256)
* [fix] NVRTC execution backend

* [fmt] run pre-commit

* [fix] coderabbit reviews

* [test] add cuda-python to test dep

* [fix] coderabbit reviews

* [fix] CUDA 13 compatibility

* [fix] sm90

* [fix] CUDA 13 compatibility

* [fix] pre-commit

* [fix] always use cuda::std::__atomic_ref_impl

* [fix] restore to external API

* Revert "[fix] restore to external API"

This reverts commit 49bd875638.

* [fmt] use space instead tabs for py codegen

* [fix] im2col API

* [fix] revert atomic.h

* [fix] dynamic shape

* [refactor] extract common utils

* [feat] support L2 persistent map

* [fix] l2 persistent map

* [fix] pre-commit

* [fix] restore _TYPE_MAP

* [fix] pre-commit

* [fix] avoid duplicate TMA descs

* [docs] add docstring

* [fix] coderabbit

* [fix] coderabbit

* [fix] coderabbit

* [fix] coderabbit
2025-11-15 11:43:03 +08:00
5f202fe5c1 [Language] Initial version of tilelang frontend v2 (#1120)
* tilelang frontend v2

* syntax sugar: defining a local var by annotation

* [Refactor] fix type linting warning like `T.float32`

* Add tl.local_var_init for new tl.float32

* allow passing default argument as function annotation

* allow default arguments as annotation

* fix lint error

* minor fix

* [Refactor] refactor tilelang.jit and tilelang.autotune

* minor fix

* minor fix

* minor fix

* fix metal get function name

* add par_compile impl and tests

* Type consistency on tvm datatype
1. isinstance(tl.float32, tvm.DataType) == True
2. Allow `tl.float32` as function annotations
3. Allow `tl.float32` as argument to be passed to `tl.alloc` or other functions

* fix lint error

* add more warning in frontend

* update tvm version

* Minor fix on tvm_ffi annotations

* add document and examples

* fix lint error

* Simplify index calculations in example_chunk_o_bwd.py

Refactor index calculations for dg_last_fragment assignment.

* minor fix

* lint fix

---------

Co-authored-by: Lei Wang <leiwang1999@outlook.com>
Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>
2025-11-03 18:15:56 +08:00
Lei WangandXuehai Pan 10911e280f [FFI] Rebase tvm to v0.22.0 to utilize tvm-ffi (#1108)
* 3rdparty tvm bump

* bump tvm into v0.22.0

* lint fix

* rebase tvm

* Update submodule tvm to latest commit 3085bc4

* Refactor: Update configuration retrieval in CopyNode and adjust test registration in tilelang

* test fix

* add requirement

* atomic_fix

* atomic_fix

* phaseout py39

* optimize

* optimize

* lint fix

* do not clean cache

* do not clean cache

* [Minor] Minor update for Python versions and dependencies

* [Lint] fix lint for py39

* [Lint] fix lint for ROCm

* [Build][CI] Sync CI changes from upstream/sdist

* [Lint] fix lint for ROCm

* [Build][CI] Update `repair-wheel-command`

* [Minor] update abi3audit result format

* [Lint] fix lint for ROCm

* [BugFix] fix build

* [Lint] fix lint for ROCm

* [BugFix] set rpath for libtvm and libtvm_runtime

* [Deps] pin apache-tvm-ffi version

* [Build] set Python 3.9 Limited API for Cython target

* [Build] set Python 3.9 Limited API for Cython target

* [Deps] Restore Python 3.8 support

* [Build] use `apache-tvm-ffi`'s `libtvm_ffi`

* [BugFix] use `;` as delimiter for RPATH on macOS

* [BugFix] use `--ignore-missing-dependencies` for `delocate-wheel`

* [Build] support `sccache` if available

* [Build] add CIBW import test

* [Build][CI] enable ccache for CIBW on Linux

* [BugFix] set rpath for libtvm and libtvm_runtime

* Revert "[Build][CI] enable ccache for CIBW on Linux"

This reverts commit cd9ab57bb5.

* [CI] fix perfbench bot

* [BugFix] use Python 3.9 to build wheel

* [Minor] update perfbench bot envs

* [BugFix] fix CIBW environment on Linux

* [CI] skip import test on CentOS 7

* [CI] use Python urllib to download file instead of Wget

---------

Co-authored-by: Xuehai Pan <XuehaiPan@pku.edu.cn>
2025-10-31 13:17:50 +08:00
Wenhao Xie 50e789dd14 [Feature] Support None type as input for T.ptr and T.Tensor (#1114)
* [Feature] Support None type as input for T.ptr and T.Tensor

* lint

* lint

* lint

* lint fix
2025-10-23 20:19:14 +08:00
Lei Wang a7730272e4 [Language] Recommend using T.dynamic instead of T.symbolic (#1076)
* recommend using T.dynamic instead of T.symbolic

* lint fix

* lint fix
2025-10-20 19:48:31 +08:00
Lei Wang 8c8d8ca23b [Refactor] Refactor jit to _JitImplementation to support @tilelang.jit (#502)
* [Refactor] Rename `jit` class to `_JitImplementation` and improve debug path handling

* Refactored the `jit` class to `_JitImplementation` for clarity and encapsulation.
* Enhanced handling of `debug_root_path` to ensure it is correctly set as an absolute path when provided.
* Updated the public `jit` function to serve as a decorator interface, allowing for both default and configured usage.
* Added validation to ensure input tensors are contiguous in the Cython wrapper, improving error handling.

* [Refactor] Improve formatting and handling in `_JitImplementation` and `jit` function

* Refactored the `_JitImplementation` class to enhance readability by adjusting comment formatting and consolidating conditions for setting `debug_root_path`.
* Updated the `jit` function signature for better alignment and clarity in parameter definitions.
* Ensured consistent spacing and comments throughout the code for improved maintainability.

* [Refactor] Update GEMM test parameters for performance optimization

* Set num_stages to 0 and adjusted matrix dimensions in the GEMM test function to enhance performance and consistency across tests in test_tilelang_jit_gemm.py.
* Reduced the number of threads used in the test to align with the updated configuration, improving overall test efficiency.

* [Refactor] Enhance buffer error logging in layout inference

* Updated the warning message in layout inference to provide clearer context when a buffer cannot be inferred due to its absence in the use list. This change improves the clarity of error reporting during layout inference operations.
* Refactored tensor handling in the Cython wrapper to ensure input tensors are checked for contiguity before processing, enhancing error handling and robustness in tensor management.

* bugfix
2025-05-20 11:47:20 +08:00
Lei Wang 25a50f1a37 [Refactor] refactor tilelang.jit to support a faster and more flexible kernel cache (#501)
* [Refactor] Update JIT kernel functions and streamline GEMM tests

* Renamed and refactored matmul and run_gemm functions to matmul_kernel_jit and run_gemm_kernel_jit for clarity.
* Removed redundant JIT decorator from the matmul function, ensuring it is applied only to the kernel function.
* Updated test function names to reflect changes in the kernel functions, enhancing consistency and readability.
* Cleaned up commented-out code and unnecessary imports to improve overall code quality.

* Update main function call in GEMM test to use tilelang testing framework

* Update README and example scripts to include JIT decorator comments

* Added comments in README.md and various example scripts to indicate the use of the @tilelang.jit decorator for returning torch functions.
* Removed redundant comments that previously instructed to add the decorator, streamlining the documentation and improving clarity.

* Update GEMM test parameters for improved performance

* Set num_stages to 0 and adjusted matrix dimensions in test functions to enhance performance and consistency across GEMM tests in test_tilelang_kernel_gemm.py.
2025-05-19 02:21:11 +08:00
Lei Wang bf8a6fc142 [Refactor] Deprecated T.Buffer as arguments and rename related calls into T.Tensor (#281)
* [Refactor] Improve flash attention example and layout comparison logic

- Removed unnecessary annotation for `lse_local_split` in the flash attention example to streamline the code.
- Updated the handling of `lse_local_split` to utilize parallel processing for better performance.
- Refactored kernel compilation and profiling logic to enhance clarity and maintainability in the flash attention example.
- Added a condition in `FragmentNode::IsEqual` to handle broadcast cases, improving the robustness of layout comparisons.

* lint fix

* [Enhancement] Add support for shared memory scope in Fill operation

- Introduced handling for `shared.dyn` and `shared` memory scopes in the Fill operation.
- Implemented parallel operation and layout inference for improved performance in shared memory scenarios.
- Updated thread loop partitioning and vectorization logic to accommodate new memory scope handling.

* [Refactor] Remove deprecated decorator and enhance Cython kernel handling

- Removed the deprecated decorator from the main module and added a new implementation in the utils module for better organization.
- Introduced a pointer map in the Cython kernel adapter to manage pointer arguments, improving runtime shape resolution.
- Updated the Cython kernel wrapper to utilize the new pointer map for handling kernel arguments.
- Enhanced error checking in the tensor utility functions to ensure static shapes are enforced.
- Added a new proxy module for buffer and tensor handling, streamlining the interface for TIR programs.

* [Feature] Add matrix multiplication test and kernel implementation

- Introduced a new test file `test_tilelang_language_ptr.py` that implements a matrix multiplication function using TileLang's primitives.
- The `matmul_test` function defines a kernel for performing tile-level GEMM operations with customizable block sizes and data types.
- Added a `run_matmul` function to compile and execute the kernel, along with a test function to validate the implementation.
- Updated the `proxy.py` file to enhance type handling for buffer and tensor proxies, ensuring compatibility with TIR programs.
- Minor formatting improvements in `deprecated.py` for better readability.

* lint fix

* [Refactor] Update tensor creation in matrix multiplication test

- Replaced `T.Tensor.from_ptr` with `T.make_tensor` in `matmul_test` for improved clarity and consistency.
- Updated imports in `__init__.py` to include `make_tensor`.
- Added `make_tensor` function in `proxy.py` to streamline tensor creation from pointers.

* [Refactor] Update tensor definitions across multiple files

- Replaced instances of `T.Tensor` with updated tensor definitions in various benchmark and example files to enhance consistency and clarity.
- Adjusted tensor shapes and types in functions related to matrix multiplication, attention mechanisms, and other operations.
- Improved documentation in README and example files to reflect changes in tensor usage.

* lint fix

* [Refactor] Update tensor types in attention and matrix multiplication examples

- Replaced instances of `T.Tensor` with `T.SharedTensor` and `T.FragmentTensor` in various attention and matrix multiplication functions to improve consistency and clarity.
- Adjusted tensor definitions in benchmark and example files to align with the new tensor types.
- Enhanced the overall structure and readability of the code by standardizing tensor usage across multiple files.

* lint fix

* [Refactor] Update tensor types in GEMM example and test files

- Replaced instances of `T.Tensor` with `T.LocalTensor` and `T.Buffer` in the GEMM example and related test functions to improve consistency and clarity.
- Enhanced the overall structure of the code by standardizing tensor usage across multiple files, aligning with recent updates in tensor definitions.

* [Refactor] Update tensor usage in customize.py

- Replaced instances of `T.Tensor` with `T.Buffer` in the `reshape` and `view` functions to enhance consistency with recent tensor definitions.
- Improved code clarity by standardizing buffer usage across the file.

* [Refactor] Update tensor types in test_tilelang_transform_annotate_device_regions.py

- Replaced instances of `T.Tensor` with `T.Buffer` in the `before` and `expected` methods of the `TestAnnotateThreadExtent` and `TestAnnotateDeviceScope` classes to enhance consistency with recent tensor definitions.
- Improved code clarity by standardizing buffer usage across the test file.

* [Refactor] Update tensor types to SharedBuffer and FragmentBuffer

- Replaced instances of `T.SharedTensor` and `T.FragmentTensor` with `T.SharedBuffer` and `T.FragmentBuffer` across multiple benchmark, example, and test files to enhance consistency with recent tensor definitions.
- Improved code clarity and structure by standardizing buffer usage in attention and matrix multiplication functions.

* [Refactor] Introduce Tensor alias for Buffer in proxy.py

- Added a new alias `Tensor` for `Buffer` in `proxy.py` to facilitate JIT compilation, ensuring that inputs and outputs are mapped with `torch.Tensor`.
- This change enhances clarity and consistency in tensor usage across the codebase.
2025-03-26 13:42:52 +08:00
Lei Wang f47b43c503 [Refactor] Refactor CUDA post-processing callback registration in TileLang (#259)
* Add GPU kernel for 2D continuous cumulative sum in TileLang example

- Introduced a new example script `example_tilelang_cumsum.py` that generates a GPU kernel for 2D continuous cumulative sum.
- Implemented functions to handle kernel configuration, memory allocation, and inclusive scan operations.
- Added a main execution block to demonstrate the kernel's functionality using PyTorch for tensor operations.
- Enhanced the example with error handling for power-of-two configurations and validation of results against PyTorch's built-in cumulative sum function.

* Refactor TileLang examples and enhance kernel compilation

- Updated `example_tilelang_cumsum.py` to improve GPU kernel generation for 2D continuous cumulative sum, including better parameter handling and error checking.
- Refactored `example_mha_bwd.py` to enhance kernel compilation readability and maintainability.
- Modified `kernel_cache.py` to prevent saving kernels to disk when using the DLPack backend, ensuring proper cache management.
- Added `get_block_bindings` function to `kernel.py` for improved access to block bindings in kernel launch frames.
- Cleaned up import statements in `__init__.py` for better organization and clarity.

* Enhance GPU kernel for 2D continuous cumulative sum in TileLang example

- Added additional spacing for improved readability in `example_tilelang_cumsum.py`.
- Refined kernel structure to enhance clarity and maintainability during GPU kernel generation for cumulative sum operations.

* Refactor CUDA post-processing callback registration in TileLang

- Introduced a new decorator `register_cuda_postproc_callback` for registering CUDA post-processing functions, enhancing usability and flexibility.
- Updated existing callback implementations to utilize the new decorator, improving code clarity and maintainability.
- Added debug prints to the CUDA code generation process for better traceability during development.
- Refactored the `OptimizeForTarget` function to streamline conditional statement handling in the pipeline transformation.
- Cleaned up the `inject_pipeline.cc` file by removing redundant code related to statement grouping and condition handling.

* lint fix

* Enhance BlockSparse GEMM Example with Autotuning and Configurable Parameters

- Added argument parsing to allow dynamic configuration of matrix dimensions and sparsity ratio.
- Implemented a function to generate various kernel configurations for autotuning.
- Refactored the main execution block to support both autotuned and default configurations.
- Improved the block mask generation to accommodate specified sparsity levels.
- Updated the kernel compilation process to utilize the new configurations and ensure accurate results verification.
2025-03-22 17:19:43 +08:00
yyttt6 316d3b97f1 add autotune to example_gemm.py (#252)
* add autotune to example_gemm.py

* add autotune to example_gemm.py

* add autotune to example_gemm.py

* add autotune to example_gemm.py
2025-03-21 12:54:06 +08:00
Lei Wang 2d0c416978 [Enhancement] Support float variable as arguments (#250)
* [Enhancement] Add matrix multiplication functions for integer and float variables in Cython JIT

- Introduced `matmul_int_variable` and `matmul_float_variable` functions to support matrix multiplication with dynamic shapes and additional parameters.
- Implemented corresponding `run_matmul_int_variable` and `run_matmul_float_variable` functions for testing.
- Updated test cases to validate the new matrix multiplication implementations.
- Enhanced error handling in library initialization and compilation processes across various modules.
- Improved dynamic memory handling in CUDA kernel initialization to provide better error reporting.

* lint fix

* optimize
2025-03-20 23:27:53 +08:00
Lei Wang f2e99180d5 [Refactor] Phaseout LLVM Dependency by Making it Optional (#247)
* remove llvm build

* [Refactor] Update kernel compilation and profiling in examples

- Replaced `tilelang.lower` with `tilelang.compile` in multiple example scripts to streamline kernel compilation.
- Updated profiling calls to utilize the new `get_profiler` method, enhancing performance measurement consistency.
- Adjusted assertions and benchmarking methods to align with the new profiling structure across various examples, ensuring correctness and clarity in performance evaluations.

* lint fix

* License Update

* [Refactor] Improve code formatting and documentation in CUDA header and HIP runtime files

- Adjusted formatting in `cuda.h` for better readability, including alignment of comments and struct fields.
- Cleaned up whitespace and improved comment clarity in `rt_mod_hip.cc` to enhance code maintainability.

* [Refactor] Enhance formatting and clarity in CUDA header and HIP runtime files

- Improved comment alignment and readability in `cuda.h`.
- Cleaned up whitespace and formatting in `rt_mod_hip.cc` to enhance maintainability.

* lint fix

* lint fix

* lint fix

* lint fix

* fix

* License update

* [Enhancement] Update JITKernel to use artifact for kernel source

- Assigned the generated artifact to `self.artifact` for better management.
- Updated kernel source references to use `artifact.kernel_source` for consistency in execution backend handling.

* lint fix

* Add @tilelang.testing.requires_llvm decorator to vectorization tests

* Enhance setup.py and env.py for library management

- Added functionality to remove original files after copying in CMakeBuild.
- Updated TVM_LIBRARY_PATH in env.py to include the PyPI build library path for better integration.

* Refactor TVM_LIBRARY_PATH assignment for improved readability in env.py

* Refactor CMakeBuild file handling in setup.py

- Added a check to ensure the target library directory exists before copying .so files.
- Improved the logic for creating the target directory and copying files to enhance robustness.

* bugfix

* Rename BuildTLDebug to BuildTileLangCUDAWithoutCompile and update registration. Add @tilelang.testing.requires_llvm decorator to multiple tests for LLVM requirement.

* lint fix

* Enhance TileLang code generation by adding support for device code generation without compilation. Updated `host_codegen` and `device_codegen` functions to include new transformations and registration for `tilelang_hip_without_compile`. Refactored JIT kernel adapters to accommodate host and device modules, improving overall integration and flexibility.

* lint fix

* Add support for C target in device code generation

- Updated `device_codegen_without_compile` to include handling for the C target by registering the `tilelang_cpp` function.

* [Enhancement] Implement auto-clear cache feature based on environment variable

* Added TILELANG_CLEAR_CACHE environment variable to control cache clearing.
* Updated CI workflow to set TILELANG_CLEAR_CACHE during testing.
* Modified cache initialization to clear cache if TILELANG_CLEAR_CACHE is set to true.

* [Refactor] Update kernel invocation and import paths in tests and cache

* Changed kernel invocation in `test_tilelang_kernel_dequantize_gemm.py` to return the result.
* Updated import statements in `test_tilelang_kernel_int4_gemm_mma.py` to use `bitblas` instead of `tilelang`.
* Refactored paths for artifact and parameters in `kernel_cache.py` for better maintainability.

* [Refactor] Clean up whitespace and improve code formatting in kernel_cache.py

* Removed unnecessary blank lines and adjusted spacing for better readability in the KernelCache class.
* Enhanced overall code formatting to align with project standards.

* [Enhancement] Add bfloat16 test case and improve kernel caching logic

* Introduced a new test case for bfloat16 matrix multiplication in `test_tilelang_kernel_gemm_mma_intrinsic.py`.
* Updated `KernelCache` to handle multiple kernel source files and improve error handling during saving and loading.
* Refactored `JITKernel` to support instantiation from a database, enhancing flexibility in kernel management.
* Adjusted `CtypesKernelAdapter` and `CythonKernelAdapter` to utilize the new kernel loading mechanism from the database.
* Improved code formatting and readability across several files.

* lint fix

* Update bfloat16 matrix multiplication test case to use larger dimensions for improved coverage
2025-03-20 20:08:20 +08:00
Wenhao Xie 71537ba567 [BugFix] Fix bug of mismatching dtype in testing (#245)
* [Typo] Fix formatting in installation instructions in README.md

* [Enhancement] Improve CUDA path detection and update configuration handling

* fix typo

* remove IS_WINDOWS constant

* lint fix

* Improve error messages for CUDA detection failure

* lint fix

* lint fix

* Fix .gitignore to correctly include venv directory

* [Doc] Add instructions for installing nightly version of TileLang

* update installation instructions

* update install instruction

* fix bug of mismatching dtype in testing and set the default value of check_dtype in torch_assert_close to true

* lint fix

* fix bug

* use map_torch_type
2025-03-19 13:43:47 +08:00
Lei Wang 3de9f13cb1 [Refactor] Introduce KernelParam integration across modules (#223)
* [Refactor] Update KernelParam integration across modules

- Replaced instances of TensorType with KernelParam in various modules to standardize parameter handling.
- Updated JITKernel, BaseKernelAdapter, and CythonKernelAdapter to utilize KernelParam for improved type consistency.
- Enhanced Profiler class to include KernelParam in its parameters, ensuring better integration with the new parameter structure.
- Adjusted tensor handling in utility functions to accommodate the new KernelParam type, improving overall code clarity and maintainability.
- Updated copyright headers to reflect the correct organization.

* [Refactor] Clean up whitespace in kernel, profiler, and tensor modules

- Added blank lines for improved readability in kernel.py, __init__.py, and tensor.py.
- Enhanced code clarity by ensuring consistent formatting across these modules.

* [Enhancement] Add detailed docstrings to KernelParam and Profiler classes

- Enhanced KernelParam class with comprehensive docstrings for better understanding of its purpose and methods.
- Updated Profiler class to include detailed docstrings for its attributes and methods, improving code documentation and usability.
- Removed unused do_bench function to streamline the profiler module and improve clarity.

* [Refactor] Update type hints in do_bench function and clean up whitespace in profiler module

- Changed type hints for grad_to_none and quantiles parameters in do_bench function to use Optional for better clarity.
- Added a blank line in __init__.py for improved readability and consistency in the profiler module.

* [Refactor] Update type hint in do_bench function for consistency

- Changed the return type hint in the do_bench function from a union type to a more explicit List type for better clarity and consistency in type annotations.

* [Refactor] Update return type hint in do_bench function for clarity

- Changed the return type hint in the do_bench function from a union type to Union[float, List[float]] for improved clarity and consistency in type annotations.

* [Enhancement] Add func property to Profiler class for adapter access

- Introduced a new property `func` in the Profiler class to provide access to the adapter, ensuring that the adapter is set before retrieval. This enhancement improves the usability of the Profiler class by allowing easier access to the adapter functionality.

* [Refactor] Update kernel compilation and profiling in tests

- Replaced instances of `TL.lower` and `TL.Profiler` with `tilelang.compile` and the new profiler interface across multiple test files.
- Enhanced the kernel compilation process to utilize the updated API, improving consistency and maintainability in the testing framework.
- Updated assertions to use the new profiler methods for better clarity and functionality in performance testing.

* [Refactor] Simplify kernel invocation and remove unused parameters in tests

- Updated the kernel invocation in `test_tilelang_dynamic_symbolic.py` to directly assign the result to `C`, improving clarity.
- Removed the `execution_backend` parameter from `tilelang.compile` calls in `test_tilelang_jit_callback.py` and `test_tilelang_jit_gemm.py` for consistency with the updated API.
- Commented out the call to `tilelang.testing.main()` in `test_tilelang_jit_callback.py` and replaced it with a direct call to `test_gemm_jit_kernel()` to streamline test execution.
- Adjusted the dtype mapping in `TorchDLPackKernelAdapter` to use the parameter's dtype directly, enhancing code simplicity.

* [Refactor] Remove unused imports in test files for cleaner code

- Eliminated unnecessary imports of `tilelang` as `TL` in various test files to enhance code clarity and maintainability.
- Updated multiple test files to streamline the codebase and reduce potential confusion from unused references.

* [Refactor] Simplify kernel invocation in tilelang kernel test

- Updated the kernel invocation in `test_tilelang_kernel_bf16_gemm_mma.py` to directly assign the result to `C`, enhancing code clarity and consistency with recent changes in the API.

* [Refactor] Simplify kernel invocation in tilelang kernel tests

- Updated kernel invocations in multiple test files to directly assign the result to `C`, improving code clarity and consistency with the updated API.
- Removed unnecessary initialization of `C` as a zero tensor, streamlining the code further.

* [Refactor] Update kernel invocation in tilelang transform tests

- Replaced the use of `TL.Profiler` with `tilelang.compile` in `test_tilelang_transform_simplify.py`, enhancing code clarity and consistency with the updated API.
- Streamlined the kernel invocation process by directly assigning the result to `C`, improving readability and maintainability of the test code.
2025-03-16 15:07:54 +08:00
Lei Wang d34601ab88 [Bugfix] Add dynamic shape support with out_idx in Cython JIT kernel compilation (#185)
* Optimize CMake build process with dynamic job count calculation

- Modify build_csrc function to use 90% of available CPU cores
- Ensure at least one job is used during compilation
- Improve build performance by dynamically adjusting parallel job count

* Optimize build_csrc function with multiprocessing module

- Replace os.cpu_count() with multiprocessing.cpu_count()
- Maintain existing 90% CPU utilization logic
- Improve CPU core count calculation for build process

* Add dynamic shape support with out_idx in Cython JIT kernel compilation

- Implement `run_cython_dynamic_shape_with_out_idx` function in test_tilelang_jit_gemm_cython.py
- Update Cython wrapper to handle dynamic symbolic shapes during tensor allocation
- Add support for resolving dynamic shape dimensions using input tensor references
- Enhance flexibility of JIT kernel compilation with symbolic shape handling

* Enhance error reporting for dynamic symbolic shape resolution in Cython JIT kernel

- Add detailed error message when a dynamic symbolic dimension is not found in dynamic_symbolic_map
- Improve debugging by providing context about missing symbolic dimensions
- Maintain existing dynamic shape resolution logic
2025-03-11 15:08:16 +08:00
Lei Wang 3471904f68 [JIT] Support Cython jit and make cython a default execution backend (#102)
* [Feature] Add CTypes JIT kernel support for dynamic shapes and multi-stream execution

- Enhance CtypesKernelAdapter to handle dynamic symbolic shapes
- Add support for multi-stream kernel execution in CTypes backend
- Implement dynamic shape handling in test_tilelang_jit_gemm_ctypes.py
- Add symbolic shape utility function in tilelang.language
- Update profiler to improve flexibility in benchmark selection

* Remove redundant thread binding in GEMM kernel implementations

- Remove unnecessary `thread_binding` line in GEMM kernel functions
- Clean up code in `examples/gemm/README.md` and `testing/python/kernel/test_tilelang_kernel_int4_gemm_mma.py`
- Enhance code readability by removing redundant thread binding annotation

* Fix indentation in int4 GEMM kernel test file

- Correct indentation for function calls in `test_tilelang_kernel_int4_gemm_mma.py`
- Remove extra indentation in `mma_emitter.ldmatrix_a()` and `mma_emitter.ldmatrix_b()` calls
- Improve code formatting for better readability

* [Feature] Add Cython JIT kernel support for dynamic shapes and multi-stream execution

- Implement CythonKernelAdapter to handle dynamic symbolic shapes
- Add support for multi-stream kernel execution in Cython backend
- Create comprehensive test suite for Cython GEMM kernel in test_tilelang_jit_gemm_cython.py
- Update JITKernel to include "cython" as a valid execution backend
- Add Cython-specific wrapper and library generation modules
- Update .gitignore to exclude Cython cache directory
- Modify setup.py to include Cython source files in package data

* lint fix

* [Refactor] Replace JITKernel with compile() function for kernel compilation

- Add new `compile()` function in tilelang/jit/__init__.py as a wrapper for JITKernel
- Update multiple test files and examples to use `tilelang.compile()` instead of `tilelang.JITKernel()`
- Modify kernel adapters to support optional kernel-only source retrieval
- Update `__init__.py` to import the new `compile()` function
- Improve kernel source retrieval for different execution backends

* lint fix

* remove debug print

* Add C/C++ compiler utility module and update Cython JIT kernel support

- Introduce new `tilelang/contrib/cc.py` module with cross-platform C/C++ compiler utilities
- Add functions to detect and retrieve system C/C++ compilers
- Implement cross-compilation and shared library creation support
- Update Cython JIT kernel to validate C++ compiler availability
- Modify Cython adapter to use detected C++ compiler for library generation

* Refactor float8 dtype mapping in tensor utility module

- Move float8_dtype_map inside adapt_torch2tvm function
- Simplify global scope by localizing the dtype mapping
- Maintain existing functionality for converting torch float8 tensors to TVM ndarray

* Refactor float8 dtype mapping in tensor utility module

- Move float8_dtype_map inside adapt_torch2tvm function
- Simplify global scope by localizing the dtype mapping
- Maintain existing functionality for converting torch float8 tensors to TVM ndarray

* revert

* Enhance Cython JIT adapter with Cython compiler detection

- Add `get_cython_compiler()` function to dynamically locate Cython executable
- Update Cython adapter to use detected Cython compiler instead of hardcoded command
- Raise an exception if no Cython compiler is found
- Update requirements.txt to specify minimum PyTorch version (>=2.2.0)

* Fix Cython kernel wrapper stream handling and type annotations

- Update stream parameter type to int64_t for better compatibility
- Directly use torch.cuda.current_stream().cuda_stream instead of casting
- Improve type safety and precision in Cython kernel wrapper
2025-02-21 22:44:44 +08:00
Lei Wang 7c817d515b [Feature] Add CTypes JIT kernel support (#100)
* [Feature] Add CTypes JIT kernel support for dynamic shapes and multi-stream execution

- Enhance CtypesKernelAdapter to handle dynamic symbolic shapes
- Add support for multi-stream kernel execution in CTypes backend
- Implement dynamic shape handling in test_tilelang_jit_gemm_ctypes.py
- Add symbolic shape utility function in tilelang.language
- Update profiler to improve flexibility in benchmark selection

* Remove redundant thread binding in GEMM kernel implementations

- Remove unnecessary `thread_binding` line in GEMM kernel functions
- Clean up code in `examples/gemm/README.md` and `testing/python/kernel/test_tilelang_kernel_int4_gemm_mma.py`
- Enhance code readability by removing redundant thread binding annotation

* Fix indentation in int4 GEMM kernel test file

- Correct indentation for function calls in `test_tilelang_kernel_int4_gemm_mma.py`
- Remove extra indentation in `mma_emitter.ldmatrix_a()` and `mma_emitter.ldmatrix_b()` calls
- Improve code formatting for better readability
2025-02-20 22:33:16 +08:00
Lei Wang 7cd6b3cd74 [Bugfix] Put InjectPtxAsyncCopy Pass behind ThreadSync Pass (#97)
* bump version into v0.1.0

* [Enhancement] Add custom develop command for editable installs and update .gitignore

* [Documentation] Update README to include system dependencies installation instructions

* [Build] Update setup.py to support library file copying for both release and develop modes

* [Build] Refactor library file copying logic in setup.py

* [Documentation] Remove unnecessary install section header in Installation.md

* [Build] Add tox configuration and local distribution script for multi-Python version support

* [Build] Improve git submodule update function with better error handling

* [Build] Update LLVM configuration path in ROCm installation script

* [Build] Add .tox/ to .gitignore for tox testing environment

* [Build] Add support for TVM prebuild path configuration in CMakeLists.txt

* [Cleanup] Remove unused TVM runtime error codes header

* [Cleanup] Fix TVM grid constant type reference in CUDA module

* [Cleanup] Remove unused customized_code function from IR module

* [Feature] Add TileLang thread synchronization and storage access analysis passes

* [Build] Reorder DLL search path directories for more flexible library loading

* [Refactor] Improve thread synchronization and library path handling

- Rename ThreadSync and TileLangThreadSync functions in C++ code
- Update Python docstring for ThreadSync with more detailed description
- Reorder library path detection in tilelang environment setup
- Minor comment and code cleanup in CUDA and warp specialization modules

* [Refactor] Improve thread synchronization code style and formatting

- Standardize pointer type spacing in storage_access.h and storage_access.cc
- Update whitespace and indentation in thread_storage_sync.cc
- Reorder include statements in thread_partial_sync.cc
- Minor code formatting improvements across thread synchronization files

* [Refactor] Fix global function registration for ThreadSync

- Correct global function registration to use ThreadSync instead of TileLangThreadSync
- Update TVM global registration to match recent refactoring efforts

* [Refactor] Simplify ThreadSync global function registration

- Remove unnecessary whitespace in global function registration
- Compact the TVM global registration line for ThreadSync

* [Feature] Add WebGPU code generation support in TileLang

- Implement WebGPU code generator (codegen_webgpu.cc and codegen_webgpu.h)
- Add WebGPU target support in lower.py and target.py
- Update CMakeLists.txt to include WebGPU codegen source files
- Introduce WebGPU-specific code generation for WGSL shader language

* [Refactor] Improve WebGPU code generation formatting and readability

- Enhance code formatting in codegen_webgpu.cc and codegen_webgpu.h
- Standardize pointer type spacing and indentation
- Improve line breaks and reduce line length for better readability
- Minor code style improvements in WebGPU code generation

* [Test] Add WebGPU matrix multiplication code generation test

- Implement test_webgpu_codegen.py for WebGPU matrix multiplication
- Add assert_gemm_codegen function to validate WebGPU code generation
- Include basic matrix multiplication kernel test case

* Update README with WebGPU codegen support announcement

* Support multi version pypi package build via tox

* Add support for CPU device backend with C code generation

- Introduce `is_cpu_device_backend` function to detect CPU backend with C code generation
- Modify `lower` function to handle special case of CPU device backend
- Update host and device call filtering for CPU backend
- Add conditional source code generation for C host target
- Extend JITKernel to support optional target_host parameter

* lint fix

* Enhance JIT kernel adapters with CTypes and Torch C++ backends

- Add CtypesKernelAdapter with dynamic library generation and kernel wrapping
- Implement TorchCPPKernelAdapter for CUDA kernel compilation
- Refactor BaseKernelAdapter to support more flexible initialization
- Improve error handling and argument processing in kernel adapters
- Update adapter initialization to support various execution backends

* Refactor and clean up code style in JIT CTypes adapter modules

- Apply consistent code formatting and whitespace in CTypes adapter files
- Remove unused imports and improve import organization
- Enhance readability of code in adapter, libgen, and wrapper modules
- Add missing whitespace and improve line breaks
- Minor linting and code style improvements across CTypes adapter files

* Add test for TileLang JIT GEMM with CTypes backend

- Implement comprehensive test for matrix multiplication using CTypes execution backend
- Create test functions for GEMM with float16 data type
- Add kernel source verification with custom callback
- Implement reference implementation using PyTorch for result validation
- Support various matrix multiplication configurations (transposition, block sizes)

* test fix

* Update TileLang JIT callback registration with override parameter

- Modify tilelang_callback_cuda_postproc to use @tvm.register_func(override=True)
- Ensure proper function registration with ability to replace existing implementations

* Reorder TileLang lowering passes for Hopper intrinsics and PTX async copy

- Adjust the order of LowerHopperIntrin and InjectPTXAsyncCopy passes
- Move these passes to ensure correct synchronization and device preparation

* Rebase main

* shared.dyn

* lint fix

* test fix

* Add environment variable handling for TileLang template and CUTLASS paths

- Introduce fallback logic for TL_TEMPLATE_PATH environment variable
- Add support for optional TL_CUTLASS_PATH configuration
- Include TODO comment for future environment variable renaming
2025-02-20 03:21:56 +08:00
Lei Wang 2ac51a033e [Wrap] Use a ctypes-based kernel wrapper instead of dlpack for runtime efficiency (#95)
* bump version into v0.1.0

* [Enhancement] Add custom develop command for editable installs and update .gitignore

* [Documentation] Update README to include system dependencies installation instructions

* [Build] Update setup.py to support library file copying for both release and develop modes

* [Build] Refactor library file copying logic in setup.py

* [Documentation] Remove unnecessary install section header in Installation.md

* [Build] Add tox configuration and local distribution script for multi-Python version support

* [Build] Improve git submodule update function with better error handling

* [Build] Update LLVM configuration path in ROCm installation script

* [Build] Add .tox/ to .gitignore for tox testing environment

* [Build] Add support for TVM prebuild path configuration in CMakeLists.txt

* [Cleanup] Remove unused TVM runtime error codes header

* [Cleanup] Fix TVM grid constant type reference in CUDA module

* [Cleanup] Remove unused customized_code function from IR module

* [Feature] Add TileLang thread synchronization and storage access analysis passes

* [Build] Reorder DLL search path directories for more flexible library loading

* [Refactor] Improve thread synchronization and library path handling

- Rename ThreadSync and TileLangThreadSync functions in C++ code
- Update Python docstring for ThreadSync with more detailed description
- Reorder library path detection in tilelang environment setup
- Minor comment and code cleanup in CUDA and warp specialization modules

* [Refactor] Improve thread synchronization code style and formatting

- Standardize pointer type spacing in storage_access.h and storage_access.cc
- Update whitespace and indentation in thread_storage_sync.cc
- Reorder include statements in thread_partial_sync.cc
- Minor code formatting improvements across thread synchronization files

* [Refactor] Fix global function registration for ThreadSync

- Correct global function registration to use ThreadSync instead of TileLangThreadSync
- Update TVM global registration to match recent refactoring efforts

* [Refactor] Simplify ThreadSync global function registration

- Remove unnecessary whitespace in global function registration
- Compact the TVM global registration line for ThreadSync

* [Feature] Add WebGPU code generation support in TileLang

- Implement WebGPU code generator (codegen_webgpu.cc and codegen_webgpu.h)
- Add WebGPU target support in lower.py and target.py
- Update CMakeLists.txt to include WebGPU codegen source files
- Introduce WebGPU-specific code generation for WGSL shader language

* [Refactor] Improve WebGPU code generation formatting and readability

- Enhance code formatting in codegen_webgpu.cc and codegen_webgpu.h
- Standardize pointer type spacing and indentation
- Improve line breaks and reduce line length for better readability
- Minor code style improvements in WebGPU code generation

* [Test] Add WebGPU matrix multiplication code generation test

- Implement test_webgpu_codegen.py for WebGPU matrix multiplication
- Add assert_gemm_codegen function to validate WebGPU code generation
- Include basic matrix multiplication kernel test case

* Update README with WebGPU codegen support announcement

* Support multi version pypi package build via tox

* Add support for CPU device backend with C code generation

- Introduce `is_cpu_device_backend` function to detect CPU backend with C code generation
- Modify `lower` function to handle special case of CPU device backend
- Update host and device call filtering for CPU backend
- Add conditional source code generation for C host target
- Extend JITKernel to support optional target_host parameter

* lint fix

* Enhance JIT kernel adapters with CTypes and Torch C++ backends

- Add CtypesKernelAdapter with dynamic library generation and kernel wrapping
- Implement TorchCPPKernelAdapter for CUDA kernel compilation
- Refactor BaseKernelAdapter to support more flexible initialization
- Improve error handling and argument processing in kernel adapters
- Update adapter initialization to support various execution backends

* Refactor and clean up code style in JIT CTypes adapter modules

- Apply consistent code formatting and whitespace in CTypes adapter files
- Remove unused imports and improve import organization
- Enhance readability of code in adapter, libgen, and wrapper modules
- Add missing whitespace and improve line breaks
- Minor linting and code style improvements across CTypes adapter files

* Add test for TileLang JIT GEMM with CTypes backend

- Implement comprehensive test for matrix multiplication using CTypes execution backend
- Create test functions for GEMM with float16 data type
- Add kernel source verification with custom callback
- Implement reference implementation using PyTorch for result validation
- Support various matrix multiplication configurations (transposition, block sizes)

* test fix

* Update TileLang JIT callback registration with override parameter

- Modify tilelang_callback_cuda_postproc to use @tvm.register_func(override=True)
- Ensure proper function registration with ability to replace existing implementations
2025-02-20 00:47:50 +08:00
Lei Wang 61de528855 [Dev] Support FP8 Codegen for cuda backend (#64)
* [Enhancement] Add VectorizeLoop function and update imports for compatibility

* [CI][Test] Improve test cases for vectorization and fix typos in parser comments

* lint fix

* Fix incorrect module reference for VectorizeLoop transformation

* Refactor vectorize_loop transformation by removing unused extent mutation logic

* [Enhancement] Add support for FP8 data types and global barriers in CUDA codegen

* Fix formatting in CUDA FP8 header file for consistency

* Refactor CI workflow to use 'tilelang_ci' virtual environment and update CUDA type printing for better clarity

* Update submodule 'tvm' to latest commit for improved functionality

* Refactor execution backend references from 'dl_pack' to 'dlpack' for consistency and clarity; add apply_simplify function to simplify PrimFunc or IRModule.

* Refactor CUDA code for improved readability; clean up formatting and remove unnecessary whitespace in multiple files.

* Refactor import statement in test_tilelang_kernel_dequantize_gemm.py to use 'tilelang.language' for consistency

* Add CUDA requirements to FP8 test cases and update references for clarity

* Add a blank line for improved readability in test_tilelang_kernel_fp8_gemm_mma.py

* Fix data type in reference result calculation for consistency in test_tilelang_kernel_gemm_mma_intrinsic.py

* Add CUDA requirements and FP8 test cases for matmul and gemv simulations

* Remove debug print statements and use tilelang's testing assertion for result validation in test_tilelang_kernel_gemm_mma_intrinsic.py

* Remove outdated comment regarding FP8 tests in test_tilelang_kernel_gemv_simt.py
2025-02-06 18:31:16 +08:00
Lei Wang 5e25923900 [Doc] Addd debug relevant testing and documentations (#58)
* implement jit test case

* [Dev] implement auto tune test case for matrix multiplication

* Implement test for legalize memory access and vectorized loop

* lint fix

* introduce run_once

* Refactor callback function names for consistency and improve code readability

* enhance documentations

* lint fix

* lint fix

* lint fix

* lint fix

* fix formatting issues in rt_mod_hip.cc

* add random seed initialization for deterministic testing
2025-01-26 21:11:36 +08:00
Lei Wang 38ba083b58 [Dev] Implement test case for tilelang transformations (#53)
* implement jit test case

* [Dev] implement auto tune test case for matrix multiplication

* Implement test for legalize memory access and vectorized loop

* lint fix
2025-01-25 23:53:15 +08:00
Lei Wang 39fc5a6d94 [Dev][jit] Introduce jit for kernel functions (#12)
* instruction update

* replace link with TileLang/tile-lang

* [Dev][Adapter] Implement Torch DLPack Kernel Adapter and related utilities

* lint fix

* Implement JIT Compiler Components

* Documents update

* lint fix

* update logo

* install script fix
2025-01-20 13:22:51 +08:00