Axis Dependency Tree aware code-gen and bmm example #28

yzh119 · 2021-11-24T09:49:49Z

No description provided.

yzh119 · 2021-11-24T09:51:00Z

tests/python/sparsetir/test_tir_sparse_correctness.py

+        with T.block("bmm_outer"):
+            T.block_attr({"sparse": True})
+            vb = T.axis.S(BATCH, b)
+            with T.init():


@MasterJH5574 Note that w/o these two lines we cannot identify bmm_outer as a reduction block, any ideas on what should we do to this these?

yzh119 · 2021-11-25T07:55:43Z

@MasterJH5574 It's ready for review now.

python/tvm/tir/transform/transform.py

src/tir/transforms/lower_sparse_tir.cc

MasterJH5574 · 2021-11-25T16:00:49Z

src/tir/transforms/lower_sparse_tir.cc

+        /* Add itervar to current block when
+         * - it's not used yet (not in stack) and
+         *   - it's parent axis was used in outer blocks or
+         *   - it's an iterator to a fixed axis.
+         */
+        if ((is_fixed_axis || in_stack.find(parent_axis_name) != in_stack.end()) &&
+            in_stack.find(axis_name) == in_stack.end()) {
+          loop_vars.push_back(all_loop_vars[i]);
+          axis_names.push_back(std::move(axis_name));
+          block_iters.push_back(SpIterVarToIterVar(sp_it_var, var_map));
+          iter_bindings.push_back(all_loop_vars[i]);
+          has_reduction_var |= sp_it_var->is_reduction;
+        }


I feels like when a sparse block has three iterators of type df, dv, df, the lowered TIR will have loop order of df, df, dv?

I think we should break this loop once the if isn't satisfied 🤔. It's not expected to reorder the loops in lowering.

Actually I think reorder loops at Sparse TIR stage is not useful. Because all valid reorder schedules can be performed after lowering.

src/tir/transforms/lower_sparse_tir.cc

MasterJH5574 · 2021-11-25T16:07:40Z

src/tir/transforms/lower_sparse_tir.cc

-
-    // Step 6. Create outer loops and the block binding.
-    Stmt loop = GenerateLoops(std::move(block_realize), block_iters, loop_vars);
+    // Step 5. Generate nested blocks and loops from innermost to outermost.


Haven't got enough time to review yet. I'm going to think more on this part tomorrow.

yzh119 · 2021-11-26T22:30:01Z

I'll merge this first.
We can discuss further on the generated loop order later.

* upd * upd * upd * upd * upd * upd * upd * upd * remove redundancy * fix * upd * upd

[SparseTIR] Constructors and Python Interface for `Axis` and `SparseBuffer` (#2) * add methods for Object * axis constructors * methods for SparseBuffer * put into registry * python interface [CherryPick][Intrinsic] lower_bound and upper_bound for binary search in Sparse TIR. (apache#483) (#4) * upd * upd * fix * upd * upd * upd * upd * upd * fix * upd * upd * upd * upd * upd * upd * upd * codegen-rule * upd * upd * test * upd * fix * two arguments Co-authored-by: Zihao Ye <[email protected]> Fix AxisTree (#3) * fix axis tree * upd [SparseTIR] Add SparseBufferLoad/SparseBufferStore (#5) * Add dtype for SparseBuffer * Add name for SparseBuffer. Remove `ndim` * Remove namespace sparse * Add SparseBufferLoad/Store * Add method `ndim()` [SparseTIR] Introduce SpIterVar (#6) * [SparseTIR] Introduce SpIterVar * Add conversion to PrimExpr [BugFix] Fix binary search & SpIterVar (#7) [BugFix] Add field `is_reduction` for SpIterVar (#9) * [BugFix] Add field `is_reduction` for SpIterVar * Formatting [SparseTIR] Index Lowering (#8) * Add StmtFunctor/ExprFunctor for SparseBufferStore/Load * Add basic index lowering * Finish index lowering (maybe) * Address comments * Convert CRLF to LF Frontend update, demo scripts. (#10) * Format and Buffer data structure (#1) * [SparseTIR] Constructors and Python Interface for `Axis` and `SparseBuffer` (#2) * add methods for Object * axis constructors * methods for SparseBuffer * put into registry * python interface * [CherryPick][Intrinsic] lower_bound and upper_bound for binary search in Sparse TIR. (apache#483) (#4) * upd * upd * fix * upd * upd * upd * upd * upd * fix * upd * upd * upd * upd * upd * upd * upd * codegen-rule * upd * upd * test * upd * fix * two arguments Co-authored-by: Zihao Ye <[email protected]> * Fix AxisTree (#3) * fix axis tree * upd * Format and Buffer data structure (#1) * [SparseTIR] Constructors and Python Interface for `Axis` and `SparseBuffer` (#2) * add methods for Object * axis constructors * methods for SparseBuffer * put into registry * python interface * fix axis tree * upd * Format and Buffer data structure (#1) * [SparseTIR] Constructors and Python Interface for `Axis` and `SparseBuffer` (#2) * add methods for Object * axis constructors * methods for SparseBuffer * put into registry * python interface * [CherryPick][Intrinsic] lower_bound and upper_bound for binary search in Sparse TIR. (apache#483) (#4) * upd * upd * fix * upd * upd * upd * upd * upd * fix * upd * upd * upd * upd * upd * upd * upd * codegen-rule * upd * upd * test * upd * fix * two arguments Co-authored-by: Zihao Ye <[email protected]> * Fix AxisTree (#3) * fix axis tree * upd * [SparseTIR] Add SparseBufferLoad/SparseBufferStore (#5) * Add dtype for SparseBuffer * Add name for SparseBuffer. Remove `ndim` * Remove namespace sparse * Add SparseBufferLoad/Store * Add method `ndim()` * Format and Buffer data structure (#1) * [SparseTIR] Constructors and Python Interface for `Axis` and `SparseBuffer` (#2) * add methods for Object * axis constructors * methods for SparseBuffer * put into registry * python interface * [CherryPick][Intrinsic] lower_bound and upper_bound for binary search in Sparse TIR. (apache#483) (#4) * upd * upd * fix * upd * upd * upd * upd * upd * fix * upd * upd * upd * upd * upd * upd * upd * codegen-rule * upd * upd * test * upd * fix * two arguments Co-authored-by: Zihao Ye <[email protected]> * Fix AxisTree (#3) * fix axis tree * upd * [SparseTIR] Add SparseBufferLoad/SparseBufferStore (#5) * Add dtype for SparseBuffer * Add name for SparseBuffer. Remove `ndim` * Remove namespace sparse * Add SparseBufferLoad/Store * Add method `ndim()` * [SparseTIR] Introduce SpIterVar (#6) * [SparseTIR] Introduce SpIterVar * Add conversion to PrimExpr * [BugFix] Fix binary search & SpIterVar (#7) * [BugFix] Add field `is_reduction` for SpIterVar (#9) * [BugFix] Add field `is_reduction` for SpIterVar * Formatting * upd * upd Co-authored-by: Ruihang Lai <[email protected]> [SparseTIR] SparseBlock on C++/Python side (#11) * Fix a bug in the last commit * SparseBlock on C++ & Python side [BugFix][SparseTIR] TVMScript Parser for Axis & SpIterVar (#12) * Update `cord` and `pos` * Fix `idtype` * Formatting.. * Bug fix 1 * Move new special stmts * Parser for Axis and SpIterVar * Fix context_maintainer.py [SparseTIR] Enhance SparseBlock to contain enough PrimFunc information (#13) * Enhance SparseBlock to have enough PrimFunc info * Remove `func_sparse_buffer_map_` * Don't print the map uh-huh [SparseTIR] Parser, Printer, Roundtrip (#14) * SparseBlock scope handler (part 1) * SparseBlock scope handler (part 2) * SparseBlock scope handler (part 3) * SparseBlock scope handler (fix 1) * Add SparseBufferLoad/Store on Python side * Parser for SparseBufferLoad/Store * Add SparseBlock to Python __init__ * StmtFunctor for SparseBlock * Ensure at least one dimension for SparseBuffer * Make `axis` field of SpIterVar mandatory * SparseBlock scope handler (fix 2) * Update Axis syntax by removing `name` parameter * Move to intrin.py * Add filed `from_sparse` to DenseFixedAxis * SparseTIR script printer * Roundtrip test * `update_symbol` bug fix * Fix attr visit in SparseBuffer * Define then compare in SparseBlock * Fix printer bug for SparseBuffer * Enable graph match for Axis and SparseBuffer * Complete HashReduce and EqualReduce for AxisTree and SparseBuffer * Fix typo * Rename test * Bug fix 1 * Bug fix 2 * Add more tests Move tests (#15) [SparseTIR] ReprPrinter for Axis and SpIterVar (#16) upd (#17) flatten (#18) ELL and BSR correctness test scripts (#19) [SparseTIR] SparseTIR Lowering (#20) * Fix a previous bug of sparse-fixed SpIterVar creation * Fix a previous bug in `GetDenseValue` * Refactor Collector and IndexTransformer * Construct block and loops * Fix a previous bug which rejects DV iters in collector * Update buffer map * Create root block * Fix bug of sparse-fixed SpIterVar creation * Fix bug on SpIterVar conversion (with refactor) * Fix bug when getting dependent SpIterVars * Fix bug on dependency map and index lowering * Full block read/write region * Test version 1 * Fix bug of loop order * Fix bug of batch-mm iterator ordering * Update PrimFunc args to use symbolic params * Fix bug of test "csr_element_wise" * Fix bug of index accumulation for sparse-fixed axis * Update correctness test * Test structural equality * Refactor and use Array fix nnz cols Add docstring for sparse tir lowering (#21) * add docstring * upd Add more examples part 1 (sddmm) (#22) * upd * upd * upd [SparseTIR][Schedule] SparseBlockRV, GetSparseBlock, SparseReorder (#23) * Test initialization * Fix a stupid bug of ReprPrinter * Add SparseBlockRV * Schedule: GetSparseBlock * Schedule: Reorder [SparseTIR][Schedule] GetSpIters (#24) remove hybrid script for successful compilation Add atomic intrinsic for output nonzero inference. (#25) * upd * upd Add "sparse" block attribute. (#26) Revert "remove hybrid script for successful compilation" This reverts commit eebd7c1. [SparseTIR] Hack `IsAffineBinding` check (#27) * [TensorIR][Schedule] Inherit block anotation upon creating new blocks * Fix SDDMM test * Hack IsAffineBinding for sparse blocks Axis Dependency Tree aware code-gen and bmm example (#28) * upd * upd * upd * upd * upd * upd * upd * upd * remove redundancy * fix * upd * upd Re-design Indices lowering (#29) * upd * upd * upd * upd * upd * init * format * fix * revise coding-style * format Complete indices lowering (#30) * upd * upd * upd * done * upd * passed test * upd Add more docstrings and depress warnings for new lowering algorithm. (#31) Refactor derived axis, frontend support of fusion. (#32) * upd * upd * fix Fatal bugfix and change the signature of DenseVariableAxis. (#33) Syntax simplification (#34) Change the order of generated blocks for block isolation. (#35) * upd * upd * upd Syntax of AttachAxis for BMM (#36) * upd * upd * upd [SparseTIR] Add "square sum" lowering test (#37) * Add square sum test * Remove pylint comment [BugFix] Fix offset caching in lowering (#38) * Hack compact dataflow check in a dirty way * Add two-K square sum test * Mark skipped tests * Fix offset saving in lowering Fusion syntax fix + SDDMM example. (#39) Some structure change on update offsets. (#40) [Refactor] SparseTIR Lowering (#41) * Take out methods in Scope * Refactor * Refactor "match" * Tweak scope contents * Refactor ViewIndexInAxis * Refactor Scope * SDDMM tests under implementation * Refactor block stack * Use Map for var_map * Extract NeedCreateNewBlock * Simplify SpIterVarToIterVar via GetIterExtent * Refactor NeedCreateNewBlock * Add docstring * Use "auto" correctly * Minor refactor and use some move Remove redundant analyzers (#42) Support indices lowering for attach and fuse. (#43) * upd * upd * upd Fix irregular BMM example. (#44) * upd * upd * upd * upd RGCN forward and butterfly pattern example. (#45) Fused SDDMM example. (#46) * upd * wip * fix Fix sparse reorder after refactor (#47) [Refactor] Refactor Unittest (#48) * upd * remove redundancy [Unittest] Correctness test for benchmarking scripts (#49) Bugfix and more test for axis fusion, new workload (#50) * upd * upd upd

Format and Buffer data structure (#1) [SparseTIR] Constructors and Python Interface for `Axis` and `SparseBuffer` (#2) * add methods for Object * axis constructors * methods for SparseBuffer * put into registry * python interface [CherryPick][Intrinsic] lower_bound and upper_bound for binary search in Sparse TIR. (apache#483) (#4) * upd * upd * fix * upd * upd * upd * upd * upd * fix * upd * upd * upd * upd * upd * upd * upd * codegen-rule * upd * upd * test * upd * fix * two arguments Co-authored-by: Zihao Ye <[email protected]> Fix AxisTree (#3) * fix axis tree * upd [SparseTIR] Add SparseBufferLoad/SparseBufferStore (#5) * Add dtype for SparseBuffer * Add name for SparseBuffer. Remove `ndim` * Remove namespace sparse * Add SparseBufferLoad/Store * Add method `ndim()` [SparseTIR] Introduce SpIterVar (#6) * [SparseTIR] Introduce SpIterVar * Add conversion to PrimExpr [BugFix] Fix binary search & SpIterVar (#7) [BugFix] Add field `is_reduction` for SpIterVar (#9) * [BugFix] Add field `is_reduction` for SpIterVar * Formatting [SparseTIR] Index Lowering (#8) * Add StmtFunctor/ExprFunctor for SparseBufferStore/Load * Add basic index lowering * Finish index lowering (maybe) * Address comments * Convert CRLF to LF Frontend update, demo scripts. (#10) * Format and Buffer data structure (#1) * [SparseTIR] Constructors and Python Interface for `Axis` and `SparseBuffer` (#2) * add methods for Object * axis constructors * methods for SparseBuffer * put into registry * python interface * [CherryPick][Intrinsic] lower_bound and upper_bound for binary search in Sparse TIR. (apache#483) (#4) * upd * upd * fix * upd * upd * upd * upd * upd * fix * upd * upd * upd * upd * upd * upd * upd * codegen-rule * upd * upd * test * upd * fix * two arguments Co-authored-by: Zihao Ye <[email protected]> * Fix AxisTree (#3) * fix axis tree * upd * Format and Buffer data structure (#1) * [SparseTIR] Constructors and Python Interface for `Axis` and `SparseBuffer` (#2) * add methods for Object * axis constructors * methods for SparseBuffer * put into registry * python interface * fix axis tree * upd * Format and Buffer data structure (#1) * [SparseTIR] Constructors and Python Interface for `Axis` and `SparseBuffer` (#2) * add methods for Object * axis constructors * methods for SparseBuffer * put into registry * python interface * [CherryPick][Intrinsic] lower_bound and upper_bound for binary search in Sparse TIR. (apache#483) (#4) * upd * upd * fix * upd * upd * upd * upd * upd * fix * upd * upd * upd * upd * upd * upd * upd * codegen-rule * upd * upd * test * upd * fix * two arguments Co-authored-by: Zihao Ye <[email protected]> * Fix AxisTree (#3) * fix axis tree * upd * [SparseTIR] Add SparseBufferLoad/SparseBufferStore (#5) * Add dtype for SparseBuffer * Add name for SparseBuffer. Remove `ndim` * Remove namespace sparse * Add SparseBufferLoad/Store * Add method `ndim()` * Format and Buffer data structure (#1) * [SparseTIR] Constructors and Python Interface for `Axis` and `SparseBuffer` (#2) * add methods for Object * axis constructors * methods for SparseBuffer * put into registry * python interface * [CherryPick][Intrinsic] lower_bound and upper_bound for binary search in Sparse TIR. (apache#483) (#4) * upd * upd * fix * upd * upd * upd * upd * upd * fix * upd * upd * upd * upd * upd * upd * upd * codegen-rule * upd * upd * test * upd * fix * two arguments Co-authored-by: Zihao Ye <[email protected]> * Fix AxisTree (#3) * fix axis tree * upd * [SparseTIR] Add SparseBufferLoad/SparseBufferStore (#5) * Add dtype for SparseBuffer * Add name for SparseBuffer. Remove `ndim` * Remove namespace sparse * Add SparseBufferLoad/Store * Add method `ndim()` * [SparseTIR] Introduce SpIterVar (#6) * [SparseTIR] Introduce SpIterVar * Add conversion to PrimExpr * [BugFix] Fix binary search & SpIterVar (#7) * [BugFix] Add field `is_reduction` for SpIterVar (#9) * [BugFix] Add field `is_reduction` for SpIterVar * Formatting * upd * upd Co-authored-by: Ruihang Lai <[email protected]> [SparseTIR] SparseBlock on C++/Python side (#11) * Fix a bug in the last commit * SparseBlock on C++ & Python side [BugFix][SparseTIR] TVMScript Parser for Axis & SpIterVar (#12) * Update `cord` and `pos` * Fix `idtype` * Formatting.. * Bug fix 1 * Move new special stmts * Parser for Axis and SpIterVar * Fix context_maintainer.py [SparseTIR] Enhance SparseBlock to contain enough PrimFunc information (#13) * Enhance SparseBlock to have enough PrimFunc info * Remove `func_sparse_buffer_map_` * Don't print the map uh-huh [SparseTIR] Parser, Printer, Roundtrip (#14) * SparseBlock scope handler (part 1) * SparseBlock scope handler (part 2) * SparseBlock scope handler (part 3) * SparseBlock scope handler (fix 1) * Add SparseBufferLoad/Store on Python side * Parser for SparseBufferLoad/Store * Add SparseBlock to Python __init__ * StmtFunctor for SparseBlock * Ensure at least one dimension for SparseBuffer * Make `axis` field of SpIterVar mandatory * SparseBlock scope handler (fix 2) * Update Axis syntax by removing `name` parameter * Move to intrin.py * Add filed `from_sparse` to DenseFixedAxis * SparseTIR script printer * Roundtrip test * `update_symbol` bug fix * Fix attr visit in SparseBuffer * Define then compare in SparseBlock * Fix printer bug for SparseBuffer * Enable graph match for Axis and SparseBuffer * Complete HashReduce and EqualReduce for AxisTree and SparseBuffer * Fix typo * Rename test * Bug fix 1 * Bug fix 2 * Add more tests Move tests (#15) [SparseTIR] ReprPrinter for Axis and SpIterVar (#16) upd (#17) flatten (#18) ELL and BSR correctness test scripts (#19) [SparseTIR] SparseTIR Lowering (#20) * Fix a previous bug of sparse-fixed SpIterVar creation * Fix a previous bug in `GetDenseValue` * Refactor Collector and IndexTransformer * Construct block and loops * Fix a previous bug which rejects DV iters in collector * Update buffer map * Create root block * Fix bug of sparse-fixed SpIterVar creation * Fix bug on SpIterVar conversion (with refactor) * Fix bug when getting dependent SpIterVars * Fix bug on dependency map and index lowering * Full block read/write region * Test version 1 * Fix bug of loop order * Fix bug of batch-mm iterator ordering * Update PrimFunc args to use symbolic params * Fix bug of test "csr_element_wise" * Fix bug of index accumulation for sparse-fixed axis * Update correctness test * Test structural equality * Refactor and use Array fix nnz cols Add docstring for sparse tir lowering (#21) * add docstring * upd Add more examples part 1 (sddmm) (#22) * upd * upd * upd [SparseTIR][Schedule] SparseBlockRV, GetSparseBlock, SparseReorder (#23) * Test initialization * Fix a stupid bug of ReprPrinter * Add SparseBlockRV * Schedule: GetSparseBlock * Schedule: Reorder [SparseTIR][Schedule] GetSpIters (#24) remove hybrid script for successful compilation Add atomic intrinsic for output nonzero inference. (#25) * upd * upd Add "sparse" block attribute. (#26) Revert "remove hybrid script for successful compilation" This reverts commit eebd7c1. [SparseTIR] Hack `IsAffineBinding` check (#27) * [TensorIR][Schedule] Inherit block anotation upon creating new blocks * Fix SDDMM test * Hack IsAffineBinding for sparse blocks Axis Dependency Tree aware code-gen and bmm example (#28) * upd * upd * upd * upd * upd * upd * upd * upd * remove redundancy * fix * upd * upd Re-design Indices lowering (#29) * upd * upd * upd * upd * upd * init * format * fix * revise coding-style * format Complete indices lowering (#30) * upd * upd * upd * done * upd * passed test * upd Add more docstrings and depress warnings for new lowering algorithm. (#31) Refactor derived axis, frontend support of fusion. (#32) * upd * upd * fix Fatal bugfix and change the signature of DenseVariableAxis. (#33) Syntax simplification (#34) Change the order of generated blocks for block isolation. (#35) * upd * upd * upd Syntax of AttachAxis for BMM (#36) * upd * upd * upd [SparseTIR] Add "square sum" lowering test (#37) * Add square sum test * Remove pylint comment [BugFix] Fix offset caching in lowering (#38) * Hack compact dataflow check in a dirty way * Add two-K square sum test * Mark skipped tests * Fix offset saving in lowering Fusion syntax fix + SDDMM example. (#39) Some structure change on update offsets. (#40) [Refactor] SparseTIR Lowering (#41) * Take out methods in Scope * Refactor * Refactor "match" * Tweak scope contents * Refactor ViewIndexInAxis * Refactor Scope * SDDMM tests under implementation * Refactor block stack * Use Map for var_map * Extract NeedCreateNewBlock * Simplify SpIterVarToIterVar via GetIterExtent * Refactor NeedCreateNewBlock * Add docstring * Use "auto" correctly * Minor refactor and use some move Remove redundant analyzers (#42) Support indices lowering for attach and fuse. (#43) * upd * upd * upd Fix irregular BMM example. (#44) * upd * upd * upd * upd RGCN forward and butterfly pattern example. (#45) Fused SDDMM example. (#46) * upd * wip * fix Fix sparse reorder after refactor (#47) [Refactor] Refactor Unittest (#48) * upd * remove redundancy [Unittest] Correctness test for benchmarking scripts (#49) Bugfix and more test for axis fusion, new workload (#50) * upd * upd upd

Format and Buffer data structure (#1) [SparseTIR] Constructors and Python Interface for `Axis` and `SparseBuffer` (#2) * add methods for Object * axis constructors * methods for SparseBuffer * put into registry * python interface [CherryPick][Intrinsic] lower_bound and upper_bound for binary search in Sparse TIR. (apache#483) (#4) * upd * upd * fix * upd * upd * upd * upd * upd * fix * upd * upd * upd * upd * upd * upd * upd * codegen-rule * upd * upd * test * upd * fix * two arguments Co-authored-by: Zihao Ye <[email protected]> Fix AxisTree (#3) * fix axis tree * upd [SparseTIR] Add SparseBufferLoad/SparseBufferStore (#5) * Add dtype for SparseBuffer * Add name for SparseBuffer. Remove `ndim` * Remove namespace sparse * Add SparseBufferLoad/Store * Add method `ndim()` [SparseTIR] Introduce SpIterVar (#6) * [SparseTIR] Introduce SpIterVar * Add conversion to PrimExpr [BugFix] Fix binary search & SpIterVar (#7) [BugFix] Add field `is_reduction` for SpIterVar (#9) * [BugFix] Add field `is_reduction` for SpIterVar * Formatting [SparseTIR] Index Lowering (#8) * Add StmtFunctor/ExprFunctor for SparseBufferStore/Load * Add basic index lowering * Finish index lowering (maybe) * Address comments * Convert CRLF to LF Frontend update, demo scripts. (#10) * Format and Buffer data structure (#1) * [SparseTIR] Constructors and Python Interface for `Axis` and `SparseBuffer` (#2) * add methods for Object * axis constructors * methods for SparseBuffer * put into registry * python interface * [CherryPick][Intrinsic] lower_bound and upper_bound for binary search in Sparse TIR. (apache#483) (#4) * upd * upd * fix * upd * upd * upd * upd * upd * fix * upd * upd * upd * upd * upd * upd * upd * codegen-rule * upd * upd * test * upd * fix * two arguments Co-authored-by: Zihao Ye <[email protected]> * Fix AxisTree (#3) * fix axis tree * upd * Format and Buffer data structure (#1) * [SparseTIR] Constructors and Python Interface for `Axis` and `SparseBuffer` (#2) * add methods for Object * axis constructors * methods for SparseBuffer * put into registry * python interface * fix axis tree * upd * Format and Buffer data structure (#1) * [SparseTIR] Constructors and Python Interface for `Axis` and `SparseBuffer` (#2) * add methods for Object * axis constructors * methods for SparseBuffer * put into registry * python interface * [CherryPick][Intrinsic] lower_bound and upper_bound for binary search in Sparse TIR. (apache#483) (#4) * upd * upd * fix * upd * upd * upd * upd * upd * fix * upd * upd * upd * upd * upd * upd * upd * codegen-rule * upd * upd * test * upd * fix * two arguments Co-authored-by: Zihao Ye <[email protected]> * Fix AxisTree (#3) * fix axis tree * upd * [SparseTIR] Add SparseBufferLoad/SparseBufferStore (#5) * Add dtype for SparseBuffer * Add name for SparseBuffer. Remove `ndim` * Remove namespace sparse * Add SparseBufferLoad/Store * Add method `ndim()` * Format and Buffer data structure (#1) * [SparseTIR] Constructors and Python Interface for `Axis` and `SparseBuffer` (#2) * add methods for Object * axis constructors * methods for SparseBuffer * put into registry * python interface * [CherryPick][Intrinsic] lower_bound and upper_bound for binary search in Sparse TIR. (apache#483) (#4) * upd * upd * fix * upd * upd * upd * upd * upd * fix * upd * upd * upd * upd * upd * upd * upd * codegen-rule * upd * upd * test * upd * fix * two arguments Co-authored-by: Zihao Ye <[email protected]> * Fix AxisTree (#3) * fix axis tree * upd * [SparseTIR] Add SparseBufferLoad/SparseBufferStore (#5) * Add dtype for SparseBuffer * Add name for SparseBuffer. Remove `ndim` * Remove namespace sparse * Add SparseBufferLoad/Store * Add method `ndim()` * [SparseTIR] Introduce SpIterVar (#6) * [SparseTIR] Introduce SpIterVar * Add conversion to PrimExpr * [BugFix] Fix binary search & SpIterVar (#7) * [BugFix] Add field `is_reduction` for SpIterVar (#9) * [BugFix] Add field `is_reduction` for SpIterVar * Formatting * upd * upd Co-authored-by: Ruihang Lai <[email protected]> [SparseTIR] SparseBlock on C++/Python side (#11) * Fix a bug in the last commit * SparseBlock on C++ & Python side [BugFix][SparseTIR] TVMScript Parser for Axis & SpIterVar (#12) * Update `cord` and `pos` * Fix `idtype` * Formatting.. * Bug fix 1 * Move new special stmts * Parser for Axis and SpIterVar * Fix context_maintainer.py [SparseTIR] Enhance SparseBlock to contain enough PrimFunc information (#13) * Enhance SparseBlock to have enough PrimFunc info * Remove `func_sparse_buffer_map_` * Don't print the map uh-huh [SparseTIR] Parser, Printer, Roundtrip (#14) * SparseBlock scope handler (part 1) * SparseBlock scope handler (part 2) * SparseBlock scope handler (part 3) * SparseBlock scope handler (fix 1) * Add SparseBufferLoad/Store on Python side * Parser for SparseBufferLoad/Store * Add SparseBlock to Python __init__ * StmtFunctor for SparseBlock * Ensure at least one dimension for SparseBuffer * Make `axis` field of SpIterVar mandatory * SparseBlock scope handler (fix 2) * Update Axis syntax by removing `name` parameter * Move to intrin.py * Add filed `from_sparse` to DenseFixedAxis * SparseTIR script printer * Roundtrip test * `update_symbol` bug fix * Fix attr visit in SparseBuffer * Define then compare in SparseBlock * Fix printer bug for SparseBuffer * Enable graph match for Axis and SparseBuffer * Complete HashReduce and EqualReduce for AxisTree and SparseBuffer * Fix typo * Rename test * Bug fix 1 * Bug fix 2 * Add more tests Move tests (#15) [SparseTIR] ReprPrinter for Axis and SpIterVar (#16) upd (#17) flatten (#18) ELL and BSR correctness test scripts (#19) [SparseTIR] SparseTIR Lowering (#20) * Fix a previous bug of sparse-fixed SpIterVar creation * Fix a previous bug in `GetDenseValue` * Refactor Collector and IndexTransformer * Construct block and loops * Fix a previous bug which rejects DV iters in collector * Update buffer map * Create root block * Fix bug of sparse-fixed SpIterVar creation * Fix bug on SpIterVar conversion (with refactor) * Fix bug when getting dependent SpIterVars * Fix bug on dependency map and index lowering * Full block read/write region * Test version 1 * Fix bug of loop order * Fix bug of batch-mm iterator ordering * Update PrimFunc args to use symbolic params * Fix bug of test "csr_element_wise" * Fix bug of index accumulation for sparse-fixed axis * Update correctness test * Test structural equality * Refactor and use Array fix nnz cols Add docstring for sparse tir lowering (#21) * add docstring * upd Add more examples part 1 (sddmm) (#22) * upd * upd * upd [SparseTIR][Schedule] SparseBlockRV, GetSparseBlock, SparseReorder (#23) * Test initialization * Fix a stupid bug of ReprPrinter * Add SparseBlockRV * Schedule: GetSparseBlock * Schedule: Reorder [SparseTIR][Schedule] GetSpIters (#24) remove hybrid script for successful compilation Add atomic intrinsic for output nonzero inference. (#25) * upd * upd Add "sparse" block attribute. (#26) Revert "remove hybrid script for successful compilation" This reverts commit eebd7c1. [SparseTIR] Hack `IsAffineBinding` check (#27) * [TensorIR][Schedule] Inherit block anotation upon creating new blocks * Fix SDDMM test * Hack IsAffineBinding for sparse blocks Axis Dependency Tree aware code-gen and bmm example (#28) * upd * upd * upd * upd * upd * upd * upd * upd * remove redundancy * fix * upd * upd Re-design Indices lowering (#29) * upd * upd * upd * upd * upd * init * format * fix * revise coding-style * format Complete indices lowering (#30) * upd * upd * upd * done * upd * passed test * upd Add more docstrings and depress warnings for new lowering algorithm. (#31) Refactor derived axis, frontend support of fusion. (#32) * upd * upd * fix Fatal bugfix and change the signature of DenseVariableAxis. (#33) Syntax simplification (#34) Change the order of generated blocks for block isolation. (#35) * upd * upd * upd Syntax of AttachAxis for BMM (#36) * upd * upd * upd [SparseTIR] Add "square sum" lowering test (#37) * Add square sum test * Remove pylint comment [BugFix] Fix offset caching in lowering (#38) * Hack compact dataflow check in a dirty way * Add two-K square sum test * Mark skipped tests * Fix offset saving in lowering Fusion syntax fix + SDDMM example. (#39) Some structure change on update offsets. (#40) [Refactor] SparseTIR Lowering (#41) * Take out methods in Scope * Refactor * Refactor "match" * Tweak scope contents * Refactor ViewIndexInAxis * Refactor Scope * SDDMM tests under implementation * Refactor block stack * Use Map for var_map * Extract NeedCreateNewBlock * Simplify SpIterVarToIterVar via GetIterExtent * Refactor NeedCreateNewBlock * Add docstring * Use "auto" correctly * Minor refactor and use some move Remove redundant analyzers (#42) Support indices lowering for attach and fuse. (#43) * upd * upd * upd Fix irregular BMM example. (#44) * upd * upd * upd * upd RGCN forward and butterfly pattern example. (#45) Fused SDDMM example. (#46) * upd * wip * fix Fix sparse reorder after refactor (#47) [Refactor] Refactor Unittest (#48) * upd * remove redundancy [Unittest] Correctness test for benchmarking scripts (#49) Bugfix and more test for axis fusion, new workload (#50) * upd * upd upd upd upd upd upd upd

yzh119 added 4 commits November 23, 2021 22:23

upd

32ab233

upd

c1f9f31

upd

0ed63b0

upd

11e91f3

yzh119 commented Nov 24, 2021

View reviewed changes

yzh119 added 2 commits November 24, 2021 17:11

upd

b8fce54

upd

9295e0e

yzh119 changed the title ~~[WIP] Axis Dependency Tree aware code-gen and bmm example~~ Axis Dependency Tree aware code-gen and bmm example Nov 25, 2021

yzh119 added 4 commits November 25, 2021 00:33

upd

d865a35

upd

111e7eb

remove redundancy

24e97cd

fix

d8874ce

MasterJH5574 reviewed Nov 25, 2021

View reviewed changes

python/tvm/tir/transform/transform.py Outdated Show resolved Hide resolved

src/tir/transforms/lower_sparse_tir.cc Outdated Show resolved Hide resolved

MasterJH5574 reviewed Nov 25, 2021

View reviewed changes

src/tir/transforms/lower_sparse_tir.cc Outdated Show resolved Hide resolved

MasterJH5574 reviewed Nov 25, 2021

View reviewed changes

yzh119 added 2 commits November 25, 2021 16:07

upd

3ffa7b0

upd

982636a

yzh119 merged commit f25aa07 into sparse Nov 26, 2021

MasterJH5574 pushed a commit that referenced this pull request Dec 22, 2021

Axis Dependency Tree aware code-gen and bmm example (#28)

9f89b6e

* upd * upd * upd * upd * upd * upd * upd * upd * remove redundancy * fix * upd * upd

yzh119 added a commit that referenced this pull request Jan 21, 2022

Axis Dependency Tree aware code-gen and bmm example (#28)

f372161

* upd * upd * upd * upd * upd * upd * upd * upd * remove redundancy * fix * upd * upd

yzh119 added a commit that referenced this pull request Jan 24, 2022

Axis Dependency Tree aware code-gen and bmm example (#28)

47d90a6

* upd * upd * upd * upd * upd * upd * upd * upd * remove redundancy * fix * upd * upd

MasterJH5574 pushed a commit that referenced this pull request Jan 26, 2022

Axis Dependency Tree aware code-gen and bmm example (#28)

5257440

* upd * upd * upd * upd * upd * upd * upd * upd * remove redundancy * fix * upd * upd

yzh119 added a commit that referenced this pull request Feb 15, 2022

Axis Dependency Tree aware code-gen and bmm example (#28)

0e6ff26

* upd * upd * upd * upd * upd * upd * upd * upd * remove redundancy * fix * upd * upd

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Axis Dependency Tree aware code-gen and bmm example #28

Axis Dependency Tree aware code-gen and bmm example #28

yzh119 commented Nov 24, 2021

yzh119 Nov 24, 2021

yzh119 commented Nov 25, 2021

MasterJH5574 Nov 25, 2021

yzh119 Nov 26, 2021

MasterJH5574 Nov 25, 2021

yzh119 commented Nov 26, 2021

Axis Dependency Tree aware code-gen and bmm example #28

Axis Dependency Tree aware code-gen and bmm example #28

Conversation

yzh119 commented Nov 24, 2021

yzh119 Nov 24, 2021

Choose a reason for hiding this comment

yzh119 commented Nov 25, 2021

MasterJH5574 Nov 25, 2021

Choose a reason for hiding this comment

yzh119 Nov 26, 2021

Choose a reason for hiding this comment

MasterJH5574 Nov 25, 2021

Choose a reason for hiding this comment

yzh119 commented Nov 26, 2021