-
Notifications
You must be signed in to change notification settings - Fork 3.5k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[TensorIR][Schedule] New primitive
reorder_block_itervar (#14448)
# Motivation
Currently the `reorder` primitive only changes the loops, and block iterable variables order would not be changed.
`transform_block_layout` can change the block iterable variables, but it requires the loops outside the given block to have no branches, which limited its usage.
This schedule primitive changes the block iterable variable order directly, with API like:
```python
def reorder_block_iter_var(self, block: BlockRV, new_order: List[int]) -> None:
"""Reorder the itervars inside a given block.
Parameters
----------
block : BlockRV
The block to be transformed.
new_order : List[int]
The new block itervar order.
"""
```
where the `new_order` is a permutation of [0, 1, ..., n-1] if n is the number of itervars in the block.
# Example
Suppose we need to change the block itervar order in block "C":
```python
@T.prim_func
def matmul(A: T.Buffer[(128, 128), "float32"], B: T.Buffer[(128, 128), "float32"], C: T.Buffer[(128, 128), "float32"]) -> None:
for i, j, k in T.grid(128, 128, 128):
with T.block("C"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
with T.init():
C[vi, vj] = 0.0
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
```
after applying:
```python
sch = tir.Schedule(matmul, debug_mask="all")
C = sch.get_block("C")
sch.reorder_block_iter_var(C, [2, 1, 0])
```
the block itervar order would be changed to `vk, vj, vi`.
```python
@T.prim_func
def matmul_after_reorder_block_iter_var(A: T.Buffer[(128, 128), "float32"], B: T.Buffer[(128, 128), "float32"], C: T.Buffer[(128, 128), "float32"]):
for i, j, k in T.grid(128, 128, 128):
with T.block("C"):
vk, vj, vi = T.axis.remap("RSS", [k, j, i])
T.reads(A[vi, vk], B[vj, vk])
T.writes(C[vi, vj])
with T.init():
C[vi, vj] = T.float32(0)
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
```- Loading branch information
Showing
10 changed files
with
283 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,148 @@ | ||
| /* | ||
| * Licensed to the Apache Software Foundation (ASF) under one | ||
| * or more contributor license agreements. See the NOTICE file | ||
| * distributed with this work for additional information | ||
| * regarding copyright ownership. The ASF licenses this file | ||
| * to you under the Apache License, Version 2.0 (the | ||
| * "License"); you may not use this file except in compliance | ||
| * with the License. You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, | ||
| * software distributed under the License is distributed on an | ||
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
| * KIND, either express or implied. See the License for the | ||
| * specific language governing permissions and limitations | ||
| * under the License. | ||
| */ | ||
| #include "../utils.h" | ||
|
|
||
| namespace tvm { | ||
| namespace tir { | ||
|
|
||
| /*! | ||
| * \brief The reorder index is not a valid permutation of | ||
| * [0, 1, ..., n-1] where n is the number of block iter vars. | ||
| */ | ||
| class InvalidReorderIndex : public ScheduleError { | ||
| public: | ||
| explicit InvalidReorderIndex(IRModule mod, Block block, Array<Integer> new_order) | ||
| : mod_(mod), block_(block), new_order_(new_order) {} | ||
| IRModule mod() const final { return mod_; } | ||
| String FastErrorString() const final { | ||
| return "ScheduleError: The specified reorder indices are invalid."; | ||
| } | ||
| String DetailRenderTemplate() const final { | ||
| std::ostringstream os; | ||
| os << "The user provided block itervar index order " << new_order_ | ||
| << " is not a valid permutation of [0, 1, ..., num_block_iter_vars-1] in block {0}."; | ||
| return String(os.str()); | ||
| } | ||
| Array<ObjectRef> LocationsOfInterest() const final { return {block_}; } | ||
|
|
||
| private: | ||
| IRModule mod_; | ||
| Block block_; | ||
| Array<Integer> new_order_; | ||
| }; | ||
|
|
||
| class BlockIterVarRewriter : public StmtMutator { | ||
| public: | ||
| Map<Block, Block> block_map; | ||
| explicit BlockIterVarRewriter(const BlockNode* block_n, std::vector<int> order) | ||
| : order_(std::move(order)), block_to_rewrite(block_n) {} | ||
|
|
||
| private: | ||
| std::vector<int> order_; | ||
| const BlockNode* block_to_rewrite; | ||
| Stmt VisitStmt_(const BlockRealizeNode* op) final { | ||
| if (op->block.get() == block_to_rewrite) { | ||
| auto block_n = CopyOnWrite(op->block.get()); | ||
| Block block = op->block; | ||
| Array<IterVar> new_iter_vars; | ||
| Array<PrimExpr> new_iter_values; | ||
| for (int idx : order_) { | ||
| new_iter_vars.push_back(block->iter_vars[idx]); | ||
| new_iter_values.push_back(op->iter_values[idx]); | ||
| } | ||
| block_n->iter_vars = new_iter_vars; | ||
| Block new_block(block_n); | ||
| block_map.Set(block, new_block); | ||
| auto block_realize_n = CopyOnWrite(op); | ||
| block_realize_n->block = new_block; | ||
| block_realize_n->iter_values = new_iter_values; | ||
| return BlockRealize(block_realize_n); | ||
| } else { | ||
| return StmtMutator::VisitStmt_(op); | ||
| } | ||
| } | ||
| }; | ||
|
|
||
| void ReorderBlockIterVar(ScheduleState self, const StmtSRef& block_sref, | ||
| const Array<Integer>& new_order) { | ||
| const BlockNode* block_n = TVM_SREF_TO_BLOCK(block_sref); | ||
| std::vector<int> new_order_vec; | ||
| for (const Integer& x : new_order) { | ||
| new_order_vec.push_back(x->value); | ||
| } | ||
| // check whether new_order is valid or not; | ||
| size_t num_block_itervars = block_n->iter_vars.size(); | ||
| std::set<int> ind_set(new_order_vec.begin(), new_order_vec.end()); | ||
| bool is_full = new_order_vec.size() == num_block_itervars; | ||
| bool is_unique = (ind_set.size() == new_order_vec.size()); | ||
| bool is_within_boundary = std::all_of(new_order_vec.begin(), new_order_vec.end(), [&](int x) { | ||
| return x >= 0 && x < static_cast<int>(num_block_itervars); | ||
| }); | ||
| if (!is_full || !is_unique || !is_within_boundary) { | ||
| throw InvalidReorderIndex(self->mod, GetRef<Block>(block_n), new_order); | ||
| } | ||
|
|
||
| // find parent block | ||
| const BlockNode* parent_block_n = nullptr; | ||
| const StmtSRefNode* p = block_sref.get()->parent; | ||
| while (p != nullptr) { | ||
| if (p->stmt->IsInstance<BlockNode>()) { | ||
| parent_block_n = TVM_SREF_TO_BLOCK(GetRef<StmtSRef>(p)); | ||
| break; | ||
| } | ||
| p = p->parent; | ||
| } | ||
| const StmtSRef parent_block_sref = GetRef<StmtSRef>(p); | ||
| const Block& parent_block = GetRef<Block>(parent_block_n); | ||
|
|
||
| // rewrite block and blockrealize | ||
| BlockIterVarRewriter rewriter(block_n, std::move(new_order_vec)); | ||
| Block new_parent_block = Downcast<Block>(rewriter(parent_block)); | ||
| rewriter.block_map.Set(parent_block, new_parent_block); | ||
| self->Replace(parent_block_sref, new_parent_block, rewriter.block_map); | ||
| } | ||
|
|
||
| struct ReorderBlockIterVarTraits : public UnpackedInstTraits<ReorderBlockIterVarTraits> { | ||
| static constexpr const char* kName = "ReorderBlockIterVar"; | ||
| static constexpr bool kIsPure = false; | ||
|
|
||
| private: | ||
| static constexpr size_t kNumInputs = 2; | ||
| static constexpr size_t kNumAttrs = 0; | ||
| static constexpr size_t kNumDecisions = 0; | ||
|
|
||
| static void UnpackedApplyToSchedule(Schedule sch, BlockRV block, Array<Integer> new_order) { | ||
| sch->ReorderBlockIterVar(block, new_order); | ||
| } | ||
|
|
||
| static String UnpackedAsPython(Array<String> outputs, String block, Array<Integer> new_order) { | ||
| PythonAPICall py("reorder_block_iter_var"); | ||
| py.Input("block", block); | ||
| py.Input("new_order", new_order); | ||
| return py.Str(); | ||
| } | ||
|
|
||
| template <typename> | ||
| friend struct ::tvm::tir::UnpackedInstTraits; | ||
| }; | ||
|
|
||
| TVM_REGISTER_INST_KIND_TRAITS(ReorderBlockIterVarTraits); | ||
|
|
||
| } // namespace tir | ||
| } // namespace tvm |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,86 @@ | ||
| # Licensed to the Apache Software Foundation (ASF) under one | ||
| # or more contributor license agreements. See the NOTICE file | ||
| # distributed with this work for additional information | ||
| # regarding copyright ownership. The ASF licenses this file | ||
| # to you under the Apache License, Version 2.0 (the | ||
| # "License"); you may not use this file except in compliance | ||
| # with the License. You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, | ||
| # software distributed under the License is distributed on an | ||
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | ||
| # KIND, either express or implied. See the License for the | ||
| # specific language governing permissions and limitations | ||
| # under the License. | ||
|
|
||
| import pytest | ||
| import tvm | ||
| import tvm.testing | ||
| from tvm import tir | ||
| from tvm.script import tir as T | ||
| from tvm.tir.schedule.testing import verify_trace_roundtrip | ||
|
|
||
|
|
||
| @T.prim_func | ||
| def matmul( | ||
| A: T.Buffer((128, 128), "float32"), | ||
| B: T.Buffer((128, 128), "float32"), | ||
| C: T.Buffer((128, 128), "float32"), | ||
| ) -> None: | ||
| for i, j, k in T.grid(128, 128, 128): | ||
| with T.block("C"): | ||
| vi, vj, vk = T.axis.remap("SSR", [i, j, k]) | ||
| with T.init(): | ||
| C[vi, vj] = 0.0 | ||
| C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk] | ||
|
|
||
|
|
||
| @T.prim_func | ||
| def matmul_after_reorder_block_iter_var( | ||
| A: T.Buffer((128, 128), "float32"), | ||
| B: T.Buffer((128, 128), "float32"), | ||
| C: T.Buffer((128, 128), "float32"), | ||
| ): | ||
| for i, j, k in T.grid(128, 128, 128): | ||
| with T.block("C"): | ||
| vk, vj, vi = T.axis.remap("RSS", [k, j, i]) | ||
| T.reads(A[vi, vk], B[vj, vk]) | ||
| T.writes(C[vi, vj]) | ||
| with T.init(): | ||
| C[vi, vj] = T.float32(0) | ||
| C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk] | ||
|
|
||
|
|
||
| def test_reorder_block_iter_var(): | ||
| sch = tir.Schedule(matmul, debug_mask="all") | ||
| C = sch.get_block("C") | ||
| sch.reorder_block_iter_var(C, [2, 1, 0]) | ||
| tvm.ir.assert_structural_equal(matmul_after_reorder_block_iter_var, sch.mod["main"]) | ||
| verify_trace_roundtrip(sch=sch, mod=matmul) | ||
|
|
||
|
|
||
| def test_reorder_block_iter_var_fail_not_full(): | ||
| sch = tir.Schedule(matmul, debug_mask="all") | ||
| C = sch.get_block("C") | ||
| with pytest.raises(tvm.tir.ScheduleError): | ||
| sch.reorder_block_iter_var(C, [2, 1]) | ||
|
|
||
|
|
||
| def test_reorder_block_iter_var_fail_not_within_bound(): | ||
| sch = tir.Schedule(matmul, debug_mask="all") | ||
| C = sch.get_block("C") | ||
| with pytest.raises(tvm.tir.ScheduleError): | ||
| sch.reorder_block_iter_var(C, [-1, 3, 2]) | ||
|
|
||
|
|
||
| def test_reorder_block_iter_var_fail_not_unique(): | ||
| sch = tir.Schedule(matmul, debug_mask="all") | ||
| C = sch.get_block("C") | ||
| with pytest.raises(tvm.tir.ScheduleError): | ||
| sch.reorder_block_iter_var(C, [0, 0, 2]) | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| tvm.testing.main() |