Commit 19c85907 authored by Yu Cheng's avatar Yu Cheng Committed by LeiWang1999
Browse files

[Bugfix] Fix logic error in ReduceOp when handling CUDA architecture (#316)

* [Bugfix] Fixed the handling logic of IfThenElseNode in if_stmt_binding

* [Bugfix] Fix logic error in ReduceOp when handling CUDA architecture

- Added a check for the existence of the target attribute "arch" to ensure that there is no undefined behavior when handling the specific architecture "sm_90". This change improves the robustness and compatibility of the code.
parent 3e859e67
......@@ -158,7 +158,9 @@ Stmt ReduceOp::Lower(const LowerArgs &T, arith::Analyzer *analyzer) const {
continue;
int reducing_threads = (*extent) * (*scale);
std::stringstream ss;
if (Downcast<String>(T.target->attrs["arch"]) == "sm_90") {
bool has_arch = T.target->attrs.count("arch") > 0;
if (has_arch && Downcast<String>(T.target->attrs["arch"]) == "sm_90") {
ss << "tl::AllReduce<" << this->MakeCodegenReducer() << ", "
<< reducing_threads << ", " << (*scale) << ">::run_hopper";
} else {
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment