-
Notifications
You must be signed in to change notification settings - Fork 5.7k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[CINN] Add tile tactic and bind cuda tactic (#60534)
* [CINN] Add tile tactic * [CINN] Add bind cuda tactic
- Loading branch information
Showing
8 changed files
with
256 additions
and
9 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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,5 +1,7 @@ | ||
| core_gather_headers() | ||
|
|
||
| gather_srcs(cinnapi_src SRCS align_iter_space_tactic.cc) | ||
| gather_srcs(cinnapi_src SRCS tile_tactic.cc) | ||
| gather_srcs(cinnapi_src SRCS compute_inline_tactic.cc) | ||
| gather_srcs(cinnapi_src SRCS bind_cuda_tactic.cc) | ||
| gather_srcs(cinnapi_src SRCS arrange_storage_tactic.cc) |
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,58 @@ | ||
| // Copyright (c) 2023 CINN Authors. All Rights Reserved. | ||
| // | ||
| // Licensed 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 "paddle/cinn/ir/group_schedule/tactic/bind_cuda_tactic.h" | ||
| #include <unordered_map> | ||
| #include "paddle/cinn/ir/ir.h" | ||
|
|
||
| namespace cinn { | ||
| namespace ir { | ||
|
|
||
| void BindCudaTactic::Init(ScheduleContext* context) { context_ = context; } | ||
|
|
||
| const std::unordered_map<IterativeSpaceInfo::AxisType, std::string> | ||
| axis_type2bind_info = { | ||
| {IterativeSpaceInfo::AxisType::kCudaBlockX, "blockIdx.x"}, | ||
| {IterativeSpaceInfo::AxisType::kCudaBlockY, "blockIdx.y"}, | ||
| {IterativeSpaceInfo::AxisType::kCudaBlockZ, "blockIdx.z"}, | ||
| {IterativeSpaceInfo::AxisType::kCudaThreadX, "threadIdx.x"}, | ||
| {IterativeSpaceInfo::AxisType::kCudaThreadY, "threadIdx.y"}, | ||
| {IterativeSpaceInfo::AxisType::kCudaThreadZ, "threadIdx.z"}, | ||
| }; | ||
|
|
||
| void BindCudaTactic::Apply(ir::IRSchedule* sch, const std::string& block_id) { | ||
| std::vector<ir::Expr> loops = sch->GetLoops(block_id); | ||
| int loop_idx = 0; | ||
| for (int i = 0; | ||
| i < context_->iter_space_info.sp_space.size() && loop_idx < loops.size(); | ||
| ++i, ++loop_idx) { | ||
| const auto& axis = context_->iter_space_info.sp_space[i]; | ||
| const IterativeSpaceInfo::AxisType& axis_type = std::get<1>(axis); | ||
| if (axis_type2bind_info.count(axis_type) != 0) { | ||
| sch->Bind(loops[loop_idx], axis_type2bind_info.at(axis_type)); | ||
| } | ||
| } | ||
| for (int i = 0; | ||
| i < context_->iter_space_info.rb_space.size() && loop_idx < loops.size(); | ||
| ++i, ++loop_idx) { | ||
| const auto& axis = context_->iter_space_info.rb_space[i]; | ||
| const IterativeSpaceInfo::AxisType& axis_type = std::get<1>(axis); | ||
| if (axis_type2bind_info.count(axis_type) != 0) { | ||
| sch->Bind(loops[loop_idx], axis_type2bind_info.at(axis_type)); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| } // namespace ir | ||
| } // namespace cinn |
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,36 @@ | ||
| // Copyright (c) 2023 CINN Authors. All Rights Reserved. | ||
| // | ||
| // Licensed 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. | ||
|
|
||
| #pragma once | ||
|
|
||
| #include <string> | ||
| #include "paddle/cinn/ir/group_schedule/tactic/schedule_tactic.h" | ||
|
|
||
| namespace cinn { | ||
| namespace ir { | ||
|
|
||
| class BindCudaTactic final : public ScheduleTactic { | ||
| public: | ||
| void Init(ScheduleContext* context) override; | ||
|
|
||
| void Apply(ir::IRSchedule* sch, const std::string& block_id) override; | ||
|
|
||
| std::string TacticName() const override { return "BindCudaTactic"; } | ||
|
|
||
| private: | ||
| ScheduleContext* context_; | ||
| }; | ||
|
|
||
| } // namespace ir | ||
| } // namespace cinn |
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,84 @@ | ||
| // Copyright (c) 2023 CINN Authors. All Rights Reserved. | ||
| // | ||
| // Licensed 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 "paddle/cinn/ir/group_schedule/tactic/tile_tactic.h" | ||
| #include "paddle/cinn/ir/ir.h" | ||
|
|
||
| namespace cinn { | ||
| namespace ir { | ||
|
|
||
| void TileTactic::Init(ScheduleContext* context) { | ||
| context_ = context; | ||
| // fake strategy | ||
| auto GetFirstFactor = [](int num) { | ||
| int factor = 1; | ||
| for (int i = num - 1; i >= 1; --i) { | ||
| if (num % i == 0) { | ||
| return i; | ||
| } | ||
| } | ||
| }; | ||
|
|
||
| bool has_rb_iter = !context_->iter_space_info.rb_space.empty(); | ||
| bool has_sp_iter = !context_->iter_space_info.sp_space.empty(); | ||
| context_->iter_space_info.rb_space.clear(); | ||
| context_->iter_space_info.sp_space.clear(); | ||
|
|
||
| if (has_sp_iter) { | ||
| int sp_factor = GetFirstFactor(context_->bucket_info.sp_lower_bound); | ||
| context_->iter_space_info.sp_space.emplace_back( | ||
| ir::Expr(context_->bucket_info.sp_lower_bound / sp_factor), | ||
| IterativeSpaceInfo::AxisType::kCudaBlockX); | ||
| context_->iter_space_info.sp_space.emplace_back( | ||
| ir::Expr(sp_factor), | ||
| has_rb_iter ? IterativeSpaceInfo::AxisType::kCudaThreadY | ||
| : IterativeSpaceInfo::AxisType::kCudaThreadX); | ||
| context_->iter_space_info.sp_space.emplace_back( | ||
| ir::Expr(-1), IterativeSpaceInfo::AxisType::kSerial); | ||
| } | ||
|
|
||
| if (has_rb_iter) { | ||
| context_->iter_space_info.rb_space.emplace_back( | ||
| ir::Expr(context_->bucket_info.rb_lower_bound), | ||
| IterativeSpaceInfo::AxisType::kCudaThreadX); | ||
| context_->iter_space_info.rb_space.emplace_back( | ||
| ir::Expr(-1), IterativeSpaceInfo::AxisType::kSerial); | ||
| } | ||
| } | ||
|
|
||
| void TileTactic::Apply(ir::IRSchedule* sch, const std::string& block_id) { | ||
| std::vector<ir::Expr> loops = sch->GetLoops(block_id); | ||
| CHECK(loops.size() == 1 || loops.size() == 2) | ||
| << "All loops must be unified as sp_loop or rb_loop."; | ||
| if (loops.size() == 2) { | ||
| std::vector<ir::Expr> rb_factors; | ||
| for (const auto& axis : context_->iter_space_info.rb_space) { | ||
| rb_factors.push_back(std::get<0>(axis)); | ||
| } | ||
| sch->Split(loops[1], rb_factors); | ||
| loops = sch->GetLoops(block_id); | ||
| VLOG(6) << "after split rb loop of " << block_id << ": " | ||
| << sch->GetModule().GetExprs()[0]; | ||
| } | ||
| std::vector<ir::Expr> sp_factors; | ||
| for (const auto& axis : context_->iter_space_info.sp_space) { | ||
| sp_factors.push_back(std::get<0>(axis)); | ||
| } | ||
| sch->Split(loops[0], sp_factors); | ||
| VLOG(6) << "after split sp loop of " << block_id << ": " | ||
| << sch->GetModule().GetExprs()[0]; | ||
| } | ||
|
|
||
| } // namespace ir | ||
| } // namespace cinn |
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,36 @@ | ||
| // Copyright (c) 2023 CINN Authors. All Rights Reserved. | ||
| // | ||
| // Licensed 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. | ||
|
|
||
| #pragma once | ||
|
|
||
| #include <string> | ||
| #include "paddle/cinn/ir/group_schedule/tactic/schedule_tactic.h" | ||
|
|
||
| namespace cinn { | ||
| namespace ir { | ||
|
|
||
| class TileTactic final : public ScheduleTactic { | ||
| public: | ||
| void Init(ScheduleContext* context) override; | ||
|
|
||
| void Apply(ir::IRSchedule* sch, const std::string& block_id) override; | ||
|
|
||
| std::string TacticName() const override { return "TileTactic"; } | ||
|
|
||
| private: | ||
| ScheduleContext* context_; | ||
| }; | ||
|
|
||
| } // namespace ir | ||
| } // namespace cinn |