Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Feature] Add CTypes JIT kernel support #100

Merged
merged 3 commits into from
Feb 20, 2025

Conversation

LeiWang1999
Copy link
Contributor

This pull request introduces several changes to the tilelang project, focusing on improving the ctypes adapter and adding new benchmarking and testing functionalities. The most important changes include adding new test functions, modifying the CtypesKernelAdapter class to handle dynamic symbolic shapes, and updating the tilelang.language module to support symbolic variables.

Testing Enhancements:

  • Added new test functions test_ctypes_kernel_do_bench, test_ctypes_kernel_multi_stream, and test_ctypes_dynamic_shape to evaluate the performance and correctness of the ctypes kernel under different conditions.

Adapter Improvements:

  • Updated the CtypesKernelAdapter class to handle dynamic symbolic shapes by introducing a dynamic_symbolic_map and the _process_dynamic_symbolic method.
  • Added properties prim_func, srcpath, libpath, lib_code, and is_dynamic to the CtypesKernelAdapter class to improve code readability and maintainability.

Code Simplification:

  • Removed the TLCUDASourceWrapperWithDynamic class and integrated its functionality directly into the TLCUDASourceWrapper class, simplifying the codebase. [1] [2]

Language Module Enhancements:

  • Added a new symbolic function to the tilelang.language module to support the creation of symbolic variables, which can be used in dynamic shape handling.

Profiler Enhancements:

  • Added the determine_profiler method to the Profiler class to automatically select the appropriate profiler based on the provided function or module type.

…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 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
- 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
@LeiWang1999 LeiWang1999 merged commit 778dbd2 into tile-ai:main Feb 20, 2025
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant