SparseTIR
SparseTIR copied to clipboard
[Tracking Issue] Adding `SplitModuleByTag` pass
Pitch
Currently, we use PreservePreprocess
and RemovePreprocess
pass to filter and split a given IRModule by tag preprocessing
, however, format decomposition does not always produce preprocess blocks, some of them are post-processing blocks (e.g. when we transform the format of the output matrix. We should design a more general module split pass.
Proposed Solution
mod_preprocess, mod_compute = tir.transform.SplitModuleByTag("preprocess")(mod)
mod_compute, mod_postprocess = tir.transform.SplitModuleByTag("postprocess", reverse=True)(mod)