This pull request includes several changes to the bitblas library, focusing on improving the matrix multiplication operations and adding new scheduling capabilities. The most important changes involve updates to propagation handling, scheduler conditions, and test configurations.
Propagation Handling:
bitblas/ops/general_matmul/__init__.py: Updated the propagation handling to use TransformKind.LDMatrixTransform for boolean propagation and added a TODO comment to check device compatibility for propagation. [1][2]
This pull request includes several changes to the
bitblas
library, focusing on improving the matrix multiplication operations and adding new scheduling capabilities. The most important changes involve updates to propagation handling, scheduler conditions, and test configurations.Propagation Handling:
bitblas/ops/general_matmul/__init__.py
: Updated the propagation handling to useTransformKind.LDMatrixTransform
for boolean propagation and added a TODO comment to check device compatibility for propagation. [1] [2]Scheduler Conditions:
bitblas/ops/general_matmul/tilelang/dense/__init__.py
: Split conditions forcan_apply_fine_grain_scheduler
and added a newcan_apply_weight_propagation_scheduler
function to handle specific propagation scenarios. [1] [2]bitblas/ops/general_matmul/tilelang/dense/__init__.py
: Updated error message in the scheduler to provide more context on unsupported configurations.Scheduler Class:
bitblas/ops/general_matmul/tilelang/dense/matmul_tensorcore.py
: RefactoredMatmulWeightPropagationScheduler
to inherit fromMatmulFineGrainScheduler
and removed redundant configuration methods.Typing and Method Signatures:
bitblas/tl/base_hint.py
: Added type hints and convertedfrom_roller_hint
to a class method. [1] [2]Test Configurations:
testing/python/operators/test_general_matmul_ops_backend_tl.py
: Updated thematmul_finetune
function and tests to include thepropagate_b
parameter for more flexible testing. [1] [2] [3]