* [Enhancement] Add global memory load/store intrinsics for various vector widths (#1710)
* Introduced new intrinsics: ldg32, ldg64, ldg128, ldg256 for loading data from global memory.
* Added corresponding store intrinsics: stg32, stg64, stg128, stg256 for storing data to global memory.
* Updated built-in functions and templates to support these new operations.
* Enhanced documentation for usage examples and type checks in the language module.
* Updated CUDA code generation to utilize new intrinsics for improved performance.
* lint fix
* test fix
* Refactor buffer type annotations in global memory load/store intrinsics
Updated the type annotations for the `ldg` and `stg` functions to use `tir.Buffer` instead of the generic `Buffer` type. This change enhances type specificity and aligns with the TVM type system for better clarity and correctness in memory operations.
* lint fix
* [Enhancement] Introduce LDG/STG lowering options and transform pass
* Added configuration options to enable/disable lowering of non-predicated and predicated global memory loads/stores to ldg/stg intrinsics.
* Implemented the LowerLDGSTG transformation pass to optimize global memory access for better performance.
* Updated tests to require a minimum CUDA compute version for relevant functionality.
* Enhanced documentation for new configuration options and transformation behavior.
* lint fix
* Refactor LDG/STG intrinsic selection logic for improved readability
* Replaced switch-case statements with if-else conditions for determining the appropriate LDG operation based on byte width.
* Enhanced error handling for unsupported byte widths with a clear fatal log message.
* lint fix
* fix
* testfix
* Enhance LDGSTG pass to skip processing for CuTeDSL backend
* Added a check to skip the transformation if the target has the "cutedsl" key.
* Updated tests to require a minimum CUDA compute version of 10.0 for LDG and STG code generation tests.
* Refactor global memory load/store functions to include conditional variants
* Updated the CUDA code generation to support conditional versions of LDG and STG functions, enhancing flexibility in memory operations.
* Modified existing functions to include new conditional variants (e.g., `ldg32_conditional`, `stg32_conditional`) for better control over memory access based on predicates.
* Adjusted comments to reflect the changes in function signatures and usage.
* Update comments and input handling for LDG intrinsics to include address and optional predicate
* Modified comments for `ldg32`, `ldg64`, `ldg128`, and `ldg256` to clarify the function signatures, indicating the use of 'address' and an optional 'predicate'.
* Changed the number of inputs for these LDG functions from 2 to -1 to accommodate the new signature format.
* Add correctness verification to LDG code generation tests
* Introduced assertions to verify the correctness of generated CUDA source for `ldg32`, `ldg64`, `ldg128`, and `ldg256` functions.
* Added checks to ensure the output matches the expected reference values within specified tolerances.
* Updated test cases to enhance validation of LDG intrinsic functionality.
* test fix
* Refactor LDG and STG intrinsics to use unified naming convention
* Updated all instances of `ldg` and `stg` functions to their new names: `load_global` and `store_global`, respectively, across the codebase.
* Modified related test cases to reflect the new function names, ensuring consistency in the testing framework.
* Enhanced comments and documentation to clarify the changes in function signatures and usage.
* Update reference values in LDG256 predicated code generation test to improve accuracy
* lint fix
* Rename test function for clarity and add new test for predicated load/store global intrinsics. Update assertions to reflect new naming convention in generated kernel source.
* lint fix
* subtype cast fix
* [Refactor] Enhance variable substitution in device function generation
- Updated `wrapBodyWithHostSideAssumes` to accept a mapping of host-side variable names to device-side variables, ensuring consistent variable references during SSA conversion.
- Introduced a substitution mechanism for old parameter variables to new ones in the device function, preventing variable sharing between host and device.
- Adjusted buffer declarations to utilize new variable mappings, enhancing the robustness of the device function generation process.
* lint fix
* test fix
* lint fix
---------
Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
[Refactor] Update fp4 data structures and CUDA code generation
- Refactored `fp4_e2_2_t` structure to include methods for accessing and setting low and high bits.
- Updated `PrintVecElemLoad` and `PrintVecElemStore` methods in `codegen_cuda.cc` to reflect changes in member access for `fp4_e2_4_t` and `fp4_e2_2_t`.
- Adjusted test for `tilelang_copy_fp4` to ensure correct input dimensions and output validation based on data type consistency.
Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
* Remove RewriteUnsafeSelect Pass
* Add test for
* Add test for T.Select
* Add test for T.Select
* Add test for T.Select
* Update test_tilelang_language_select.py
---------
Co-authored-by: wangxiangwen <xiangwen.wang@deepseek.com>
Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>
* Enhance loop vectorization logic and buffer handling
- Refactored vectorization logic to utilize configuration settings for enabling/disabling vectorization.
- Introduced new checks for buffer access invariance in call nodes to improve vectorization accuracy.
- Updated GCD calculations for vector sizes to include call node constraints, enhancing overall performance.
- Added verbose logging for better insights during vectorization planning.
- Introduced tests for vectorization of infinity and bitwise logical operations to validate new functionality.
* lint fix
* fix
* test fix
* [Feature] Introduce DecoupleTypeCast pass to optimize vectorization constraints
- Added a new pass to decouple type cast vectorization constraints, allowing for optimal vectorization in mixed-precision scenarios.
- Implemented transformations for local-to-memory and memory-to-local cases, inserting intermediate local buffers as needed.
- Updated pass configuration to enable verbose output for vectorization planning.
- Enhanced loop vectorization logic to handle complex cases with multiple statements.
- Added tests to validate the new functionality and ensure correctness in various scenarios.
* Refactor verbose output in loop vectorization for improved readability
- Updated the output formatting in the VectorizePlanner to use newline characters instead of std::endl for better performance and clarity.
- Ensured consistent output style across various verbose logging statements related to buffer sizes and vectorization strategies.
* Refactor verbose output in loop vectorization for consistency
- Simplified the output formatting in the verbose logging of vector size to enhance readability by removing unnecessary line breaks.
* use eager style.
* lint fix
* Enhance loop vectorization by introducing GetBufferStrides function
- Added GetBufferStrides function to compute buffer strides when not defined, improving clarity and reusability.
- Updated VectorizeFindGlobalAccess to utilize the new function for stride calculations, simplifying the logic.
- Refactored related code to remove redundant stride calculations and improve overall readability.
- Introduced a test case to ensure no transformation occurs when different dtypes are only in if-then-else conditions, validating the new behavior.
* Refactor loop vectorization logic and enhance target support checks
- Updated the logical condition in TargetIsSm100 to use the correct logical AND operator.
- Introduced TargetSupportVectorize256 function to encapsulate vectorization support checks for specific architectures.
- Refactored VectorizeFindGlobalAccess to VectorizeFindMemoryAccess for improved clarity and added support for shared memory access checks.
- Updated vectorization logic to utilize the new support checks, enhancing overall performance and maintainability.
- Cleaned up redundant logging statements and improved code formatting for better readability.
* Update TargetSupportVectorize256 to simplify architecture check
- Modified the architecture check in TargetSupportVectorize256 to return true for architectures 100 and above, removing the upper limit of 110. This change enhances compatibility with newer architectures while maintaining existing functionality.
* lint fix
* Refactor tilelang JIT configurations and optimize GEMM calls in fused MoE example
- Updated JIT pass configurations to use enum keys for better clarity.
- Removed unused parameters `k_pack` and `coalesced_width` from the `moe_forward_tilelang_routed` function.
- Simplified GEMM calls by removing the `k_pack` argument, enhancing readability.
- Adjusted swizzle usage to remove the `enable` flag for consistency.
- Added a call to `tilelang.disable_cache()` in the main function to improve performance.
* lint fix
* Update TVM submodule and enhance vectorization logging
- Updated the TVM submodule to the latest commit for improved functionality.
- Added logging in the VectorizePlanner to provide insights during the planning phase of for loops.
- Ensured GCD calculation for vector size aligns with loop extent for better vectorization accuracy.
- Enabled verbose logging for the vectorize planner in the test configuration to aid debugging.
* lint fix
* performance fix
* fix
* fix
* fix
* fix
* bench fix
* move to vectorize_loop
* fix
* lint fix
* Refactor atomic add functions to improve type handling and reduce redundancy
- Introduced utility functions ToFloat2 and ToFloat4 for type conversion, enhancing the flexibility of atomic add operations.
- Updated AtomicAddx2 and AtomicAddx4 functions to accept generic value types, streamlining the code and improving readability.
- Removed outdated function overloads and consolidated logic for atomic addition, ensuring consistency across implementations.
- Enhanced memory order handling in atomic operations for better performance and clarity.
* add related test
* Remove atomic add vectorization implementation
- Deleted the atomicadd_vectorize.cc and atomicadd_vectorize.h files, which contained the automatic vectorization pass for atomic add operations.
- Removed references to the deleted files in atomic_add.cc and loop_partition.cc to clean up the codebase and eliminate unused dependencies.
* [Feature] Implement atomic reduction operations and enhance atomic add functionality
- Added support for atomic max and min operations with corresponding classes and methods.
- Introduced vectorization capabilities for atomic add operations based on data type and target architecture.
- Refactored atomic add intrinsic calls to improve clarity and consistency.
- Enhanced layout inference for atomic operations and integrated new utility functions for target architecture checks.
- Updated tests to validate the new atomic operations and their vectorization behavior.
* Refactor atomic operation code for improved readability and consistency
- Cleaned up whitespace and formatting in atomic add and reduce implementations for better code clarity.
- Enhanced comments in CUDA code generation for atomic operations to improve understanding.
- Updated function signatures and layout definitions for consistency across atomic operation files.
- Ensured proper alignment of code for better maintainability and readability.
* lint fix
* Enhance loop vectorization and layout handling across multiple operations
- Updated function signatures in atomic add, atomic reduce, copy, and fill operations to include layout map parameters for improved layout handling.
- Refactored vectorization logic to utilize layout maps, ensuring better performance and compatibility with various buffer layouts.
- Enhanced the LowerParallelLoop function to accept layout maps, facilitating more efficient loop transformations.
- Added checks for buffer contiguity in vectorization processes to ensure correctness when using layout maps.
- Updated tests to validate the new layout handling and vectorization behavior.
* Add atomic operation checks in legalize_safe_memory_access
- Introduced a new method to identify atomic operations within the legalize_safe_memory_access transformation.
- Updated the VisitStmt_ function to handle both CallExtern and atomic operations, ensuring recursive condition collection for these cases.
- Enhanced comments for clarity on the handling of atomic operations in the context of memory access legality.
* Add atomic operations for HIP code generation
- Implemented support for various atomic operations including atomic add, atomic load, atomic store, atomic max, and atomic min in the HIP code generation.
- Enhanced the handling of atomic operations to include optional memory order parameters.
- Improved code readability with added comments explaining the purpose of each atomic operation.
- Ensured consistency with existing atomic operation implementations in the codebase.
* Refactor AtomicOpBaseNode and related classes for improved clarity
- Changed the GetElemOpStatic method to a virtual GetElemOp method in AtomicOpBaseNode, enhancing polymorphism.
- Updated AtomicAddNode, AtomicMaxNode, and AtomicMinNode to override the new GetElemOp method.
- Removed unnecessary template parameters from AtomicOpBaseNode, simplifying the class structure.
- Cleaned up includes in atomic_reduce.cc to remove unused dependencies, improving code organization.
* Enhance argument validation for atomic operations
- Added checks in AtomicAdd, AtomicMax, and AtomicMin constructors to ensure at least two arguments are provided, improving error handling and user feedback.
- Removed the unused LowerTMA method declaration from the AtomicAdd class, streamlining the codebase.
* lint fix
* fix
* test fix
* performance fix
* remove useless comments
* Enhance CUDA code generation for BroadcastNode by implementing compile-time constant folding and runtime broadcasting for various lane configurations. Improved handling for 4-bit and 8-bit integer types, ensuring correct replication and type casting in output expressions. This update increases performance and correctness in CUDA kernel generation.
* add test
* lint fix
* fix
* support cp.reduce.async.bulk.tensor and add test
* Refactor flash attention example by removing unnecessary layout annotations
* support swizzle layout for tma reduce
* auto swizzle for non-1d tma atomic add
* upd example and test
* lint
* typo
* add constraint for test
* Refactor CUDA data type mapping by moving the to_CUtensorMapDataType function to utils.cc and utils.h, while removing redundant definitions from atomic_add.cc and copy.cc.
* lint
* [Fix] Refactor type hint extraction logic in DSLMutator for better clarity and handling of annotations
* [Refactor] Remove redundant tensor creation in loop layout tests and update kernel compilation parameters
* [Enhancement] Update global load/store functions for CUDA compatibility (#1652)
Refactor the `ld_global_256` and `st_global_256` functions to support both CUDA versions above 12.9 and earlier versions. This change ensures that 256-bit loads and stores are handled correctly across different CUDA versions, improving performance and compatibility. The implementation now uses two 128-bit loads/stores for older versions, enhancing the robustness of the codebase.
* Update comments in global load/store functions for CUDA compatibility
Clarified comments in `ld_global_256` and `st_global_256` functions to indicate that the fallback for CUDA versions below 12.9 may have performance regressions. This change enhances code readability and provides better context for developers working with different CUDA versions.
* Update submodule and enhance LetStmt handling in inject_pipeline.cc
- Updated the TVM submodule to the latest commit.
- Improved the handling of LetStmt in the inject_pipeline.cc file to account for transitive dependencies on loop variables, ensuring correct variable substitution in rewritten blocks.
- Adjusted test_tilelang_issue_1263.py to remove unnecessary jit decorator and updated the kernel compilation process with specific pass configurations.
* lint fix
* revert tvm
* remove unused test
* test fix
* Add unroll loop functionality and corresponding test case
- Introduced a new `UnrollLoop` function in the transform module to unroll loops based on various configuration options.
- Added a test case in `test_tilelang_language_unroll.py` to validate the behavior of `T.unroll` with only the extent parameter, ensuring correct kernel generation with unroll pragmas.
* Refactor unroll kernel implementation and update test case
- Changed the kernel function in `test_tilelang_language_unroll.py` to use a new `unroll_kernel` function that compiles and returns the output tensor, improving clarity and structure.
- Updated the `OptimizeForTarget` function in `phase.py` to ensure the `UnrollLoop` transformation is applied correctly, maintaining consistency in optimization phases.
* lint fix
* lint fix
- Move `tilelang/language/eager/dtypes.py` to `tilelang/language/dtypes.py`
since dtype definitions are general-purpose, not eager-specific
- Add `bytes` property to dtype (alias for `itemsize`) for convenience
- Add type hints for `bits` and `bytes` properties in TYPE_CHECKING block
- Create `tilelang/dtypes.py` for direct import via `from tilelang.dtypes import ...`
- Update all import paths accordingly
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
* Refactor: Replace @tilelang.lazy_jit with @tilelang.jit in examples and tests
Updated all instances of the @tilelang.lazy_jit decorator to @tilelang.jit in the lazyjit example notebooks and related test files. This change aligns the code with the new JIT compilation approach, enhancing consistency across the codebase. Additionally, removed the lazy_jit import from the module initialization to streamline the API.
* Enhancement: Improve layout annotation handling in T.copy and parallel loops
Updated the T.copy function to accept a new keyword-only parameter, loop_layout, allowing users to specify layout hints for the outermost parallel loop. Enhanced layout annotation handling in CopyNode and AtomicAddNode classes to ensure compatibility with SIMT operations. Added tests to validate the functionality of loop layout annotations, improving robustness and clarity in layout management for parallel loops.
* Add out_idx parameter to @jit and validate it's only used in lazy mode
* lintfix
* Refactor: centralize mode inference in JITImpl.__call__ before parse_args
* test fix
* Enhancement: Introduce JITNoBuilderError for eager mode checks in T.Kernel and T.const functions. Ensure proper error handling when Builder is unavailable, improving robustness in JIT context management.
* lint fix
* lint fix
* Enhancement: Update example_triton_sparse_gqa_decode_varlen_indice.py to improve performance metrics output by adding average FLOPS calculations. Modify layout_inference.cc to refine buffer layout inference checks by ensuring fragment buffers are considered.
* lint fix
* Enhancement: Refactor JIT implementation to centralize mode initialization and improve error handling for eager mode. Update matmul tests to utilize both Cython and TVM FFI JIT kernels, ensuring consistency in output validation.
* Refactor: Remove unused simplify_prim_func decorator from example_gemm_intrinsics.py and related test files. Enhance layout.cc to simplify integer set analysis for multi-threaded compilation, ensuring constant integer outputs. Update lower_tile_op.cc to enforce constant integer checks in layout output shapes. Improve error handling in builder.py for JIT context management.
* remove claude.md
* Refactor: Clean up code formatting in layout.cc, lower_tile_op.cc, and proxy.py for improved readability. Add missing import in test_tilelang_kernel_int4_gemm_mma.py and ensure consistent spacing in builder.py.
* Refactor: Replace logical conditions with tir_all in complex expressions within test_arith_hard.py for improved clarity and consistency. Update main execution flow to call test_hard_prove instead of tilelang.testing.main.
* Enhancement: Add @tilelang.testing.requires_cuda decorator to test functions in test_tilelang_issue_1549.py and test_tilelang_issue_1601.py to ensure CUDA requirements are explicitly stated for relevant tests.
* lint fix
* Refactor: Simplify output shape calculation in layout.cc by directly setting max value from ist.max(). Update test_valid_loop in test_tilelang_fragment_loop_checker.py to disable certain tests and ensure proper execution flow. Enhance error handling in builder.py to correctly manage PrimFunc return types.
* Refactor: Update boolean condition in test_boolop to use Or and Not for improved clarity and consistency in logical expressions.
* Refactor: Update test execution flow in test_tilelang_fragment_loop_checker.py to call tilelang.testing.main. Modify JIT implementation to handle callable functions more generically. Enhance TirTemplate class in builder.py to return an empty tuple if matcher is None, improving robustness.
* Refactor JIT function handling: Replace LazyJITFunc with JITFunc across the codebase, updating related function signatures and logic to support eager JIT compilation mode. This change enhances the clarity and functionality of JIT operations.
* Update JITImpl documentation: Clarify output tensor index attribute description and remove outdated comment regarding implicit TIR construction.
* Refactor TileLang to transition from v2 to eager language module. Update imports across multiple files to replace references to the deprecated v2 module with the new eager module, ensuring compatibility with the latest JIT compilation features. Remove obsolete v2 module files to streamline the codebase.
* Enhance JIT compilation support: Update JIT function handling to improve compatibility with eager execution. Refactor related function signatures and logic, ensuring seamless integration with the latest compilation features. This change aims to streamline JIT operations and enhance overall code clarity.
* [Enhancement] Implement hoisting of broadcast values in Let statements and update tests for validation
* lint fix
* [Enhancement] Update hoisting logic for broadcast values in Let statements and add a new test case for nested BufferStore broadcasts
* lint fix
* fix(rocm): correct index_map selection for transposed A matrix in MFMA when k_dim=4
Previously, when k_dim=4, the index_map always used non-transposed layout
for matrix A regardless of the transposed flag. This caused precision
issues for transposed GEMM operations on ROCm.
Re-enable the previously skipped test cases for trans_A=True with float dtype.
* also open int8 gemm
---------
Co-authored-by: zhutaoyu <tazhu@amd.com>
* add hoist_broadcast_values pass
* only hoist intImm or floatImm
* refactor
* add test for hoist_broadcast_values
* add a test for issue 1601
* add transform test for hoist_broadcast_value
---------
Co-authored-by: kurisu6912 <227995639+kurisu6912@users.noreply.github.com>
* 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
* [Enhancement] Improve shape and stride binding for subbyte types in ArgBinder
* Updated the shape binding logic in ArgBinder to handle subbyte types (bits < 8) by introducing a packing factor for logical shapes.
* Enhanced the handling of strides for subbyte types, ensuring correct mapping between logical and runtime strides.
* Refactored the binding process to accommodate both symbolic and constant dimensions, improving the robustness of the shape binding mechanism.
* Removed obsolete test file `test_ir_kernel_frame.py` as part of cleanup efforts.
* lint fix
* [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>
* [Enhancement] Improve layout annotation handling in parallel loops
* Updated the handling of layout annotations in parallel loops to ensure they are only attached to the outermost loop. This change clarifies the rationale behind layout management, allowing for better reasoning and transformation of nested parallel regions.
* Introduced new parameters in the `Parallel` function to accept layout annotations directly, enhancing usability and validation during layout inference.
* Refactored the layout inference logic to prioritize annotated layouts over inferred ones, improving the efficiency of layout management in parallel operations.
This update aims to strengthen the integrity of layout annotations and streamline the handling of nested parallel loops in the system.
* lint fix
* lint fix
* [Enhancement] Update test and builder for improved debugging and functionality
* Added print statements in `test_tilelang_language_let.py` to output the `main` function and the source of the compiled module, aiding in debugging.
* Modified the `bind_immutable` method in `builder.py` to include logging of the bound value type, enhancing traceability during the binding process.
* Commented out the call to `tilelang.testing.main()` in favor of directly invoking the test function, streamlining the testing workflow.
* replicate example
* [Enhancement] Add loop layout support to T.copy and improve parallel loop handling
* Updated the T.copy function to accept a new keyword-only parameter, loop_layout, allowing users to specify a layout hint for the outermost parallel loop in SIMT copies.
* Enhanced the handling of layout annotations in the CopyNode class to enforce compatibility with SIMT copy operations, ensuring that layout hints are only applied when appropriate.
* Added tests to validate the functionality of loop layout annotations in copy operations, improving the robustness of the layout management system.
* [Enhancement] Refine loop annotation handling in AtomicAdd and Copy operations
* Updated the handling of loop annotations in both AtomicAddNode and CopyNode classes to ensure that coalesced width and parallel loop layout annotations are only applied to the outermost loop.
* This change improves the clarity and correctness of loop management in SIMT operations, aligning with the recent enhancements made to layout annotation handling in parallel loops.
* refactor params of ptx_cp_aync
* Enhancement: Introduce ptx_cp_async intrinsic and improve async copy handling
- Added the `ptx_cp_async` intrinsic for asynchronous memory copy operations in TileLang, supporting both predicated and non-predicated versions.
- Updated the code generation for `ptx_cp_async` to handle offsets correctly and added inline functions for simplified usage.
- Enhanced the `inject_ptx_async_copy` transformation to support predicated vectorized copies.
- Updated example scripts to disable caching and print kernel sources for better debugging and performance insights.
* Remove tilelang.disable_cache() calls from example scripts for cleaner execution
* Refactor comments and formatting in CUDA code generation and PTX async copy transformation
- Updated comments for clarity in `codegen_cuda.cc` regarding `ptx_cp_async` arguments.
- Improved formatting in `inject_ptx_async_copy.cc` for better readability of access pointer creation and logging.
- Added a blank line in `op.py` for consistency in the `ptx_cp_async` function definition.
* Remove unused ptx_cp_async handling from CUDA code generation
- Eliminated the implementation for `ptx_cp_async` and `tl::ptx_cp_async` from `codegen_cuda.cc`, streamlining the code by removing unnecessary checks and assembly generation.
- This change focuses on cleaning up the codebase and improving maintainability by removing deprecated functionality.
* [Enhancement] Improve equality checks in layout nodes and fragment validation
* Enhanced the IsEqual method in LayoutNode and FragmentNode to include detailed comparisons of forward mappings, ensuring accurate equality checks.
* Introduced a new parameter in ProveFragmentContains to validate physical indices when checking fragment containment, improving correctness in layout validation.
* Removed obsolete test file related to layout inference.
This update aims to strengthen the integrity of layout comparisons and fragment validations in the system.
* lint fix
* [Enhancement] Improve layout inference in ParallelOpNode
* Enhanced the layout inference mechanism in ParallelOpNode to utilize two strategies: compute_loop_layout_from_buffer and PlanLoopPartition, selecting the one that minimizes replication while ensuring compatibility.
* Updated the logic to choose the best candidate layout based on replication size and containment checks.
* Refactored the HasKnownLayoutAnchor function to clarify its purpose in prioritizing buffer layouts.
* Added a new test case to validate the layout inference behavior, ensuring the correct fragments are generated in the output.
This update aims to optimize layout inference for parallel operations, improving performance and resource utilization.
* lint fix
* bug fix and refactor
* lint fix
* [Enhancement] Add C host post-processing callback and improve nullable buffer handling
- Introduced a new callback for C host code generation that allows for post-processing of generated code before it is wrapped into a CSourceModule.
- Enhanced the ArgBinder to log shape variable sources and ensure safe binding of symbolic shape variables, preventing potential segmentation faults when dealing with nullable buffers.
- Added a regression test to verify that a single buffer with a symbolic shape variable must be non-null, ensuring robustness against null inputs.
* [Refactor] Clean up formatting and improve readability in C host code generation and ArgBinder
- Adjusted indentation in the C host code generation to enhance clarity.
- Reformatted comments and code structure in ArgBinder for better readability, ensuring consistent style throughout the file.
- Minor whitespace adjustments to maintain code consistency.
* fp4 related update, require_cu13
* Enhance CUDA type conversion handling and optimize dtype management
- Updated CUDA vectorized cast functions to ensure proper handling of float16, float32, bfloat16, and float8 conversions, adding checks for bit sizes.
- Refactored dtype conversion logic in `cuda_fp4.h` to utilize `cudaRoundZero` for improved accuracy in floating-point conversions.
- Introduced a new method in `KernelParam` to convert TVM DataType to TileLang dtype.
- Adjusted argument binding logic in `arg_binder.cc` to allow for better subtype matching based on total bit counts.
- Enhanced dtype handling in `dtypes.py` to accommodate new float4_e2m1fn types and ensure compatibility with PyTorch.
This update aims to improve type safety and conversion accuracy across the codebase.
* lint fix
* lint fix
* typo fix
---------
Co-authored-by: Zhiwen Mo <zm125@ic.ac.uk>
* 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
Kuris