gemini-code-assist[bot] commented on code in PR #19857:
URL: https://github.com/apache/tvm/pull/19857#discussion_r3449011938
##########
src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc:
##########
@@ -63,12 +63,12 @@ ffi::Optional<SBlockRV> ParseAnnotate(const Schedule& sch,
const Instruction& in
}
TVM_FFI_ICHECK_EQ(inst->inputs.size(), 2);
TVM_FFI_ICHECK_EQ(inst->attrs.size(), 1);
- ffi::String ann_key = Downcast<ffi::String>(inst->attrs[0]);
+ ffi::String ann_key = (inst->attrs[0]).as_or_throw<ffi::String>();
if (ann_key != s_tir::attr::meta_schedule_cooperative_fetch) {
return std::nullopt;
}
- *vector_lane =
Downcast<IntImm>(sch->Get(Downcast<ExprRV>(inst->inputs[1])))->value;
- return Downcast<SBlockRV>(inst->inputs[0]);
+ *vector_lane =
(sch->Get((inst->inputs[1]).as_or_throw<ExprRV>())).as_or_throw<IntImm>()->value;
+ return (inst->inputs[0]).as_or_throw<SBlockRV>();
Review Comment:
### Readability: Redundant Parentheses
Similarly, the redundant parentheses around `inst->inputs[1]` and
`inst->inputs[0]` can be removed to make the code cleaner.
```suggestion
*vector_lane =
sch->Get(inst->inputs[1].as_or_throw<ExprRV>()).as_or_throw<IntImm>()->value;
return inst->inputs[0].as_or_throw<SBlockRV>();
```
##########
src/relax/backend/contrib/cudnn/codegen.cc:
##########
@@ -106,9 +106,12 @@ class cuDNNJSONSerializer : public JSONSerializer {
"kernel", /* op_type_ */
inputs, 1 /* num_outputs_ */);
const CallNode* root_call = backend::GetOpInFunction(fn,
"relax.nn.attention");
- auto q_shape =
Downcast<ShapeExpr>(Downcast<TensorType>(root_call->args[0]->ty)->shape.value());
- auto k_shape =
Downcast<ShapeExpr>(Downcast<TensorType>(root_call->args[1]->ty)->shape.value());
- auto v_shape =
Downcast<ShapeExpr>(Downcast<TensorType>(root_call->args[2]->ty)->shape.value());
+ auto q_shape =
((root_call->args[0]->ty).as_or_throw<TensorType>()->shape.value())
+ .as_or_throw<ShapeExpr>();
+ auto k_shape =
((root_call->args[1]->ty).as_or_throw<TensorType>()->shape.value())
+ .as_or_throw<ShapeExpr>();
+ auto v_shape =
((root_call->args[2]->ty).as_or_throw<TensorType>()->shape.value())
+ .as_or_throw<ShapeExpr>();
Review Comment:
### Readability: Redundant Parentheses
The newly introduced casting expressions have highly nested, redundant
parentheses which make the code difficult to read.
Since member access (`->` and `.`) and function calls (`()`) have higher
precedence than casting operations, we can safely remove the outer parentheses.
```c
auto q_shape =
root_call->args[0]->ty.as_or_throw<TensorType>()->shape.value().as_or_throw<ShapeExpr>();
auto k_shape =
root_call->args[1]->ty.as_or_throw<TensorType>()->shape.value().as_or_throw<ShapeExpr>();
auto v_shape =
root_call->args[2]->ty.as_or_throw<TensorType>()->shape.value().as_or_throw<ShapeExpr>();
```
##########
src/s_tir/meta_schedule/postproc/rewrite_cooperative_fetch.cc:
##########
@@ -41,11 +41,11 @@ ffi::Optional<int64_t> ParseThreadBinding(const Schedule&
sch, const Instruction
}
TVM_FFI_ICHECK_EQ(inst->inputs.size(), 1);
TVM_FFI_ICHECK_EQ(inst->attrs.size(), 1);
- ffi::String thread_axis = Downcast<ffi::String>(inst->attrs[0]);
+ ffi::String thread_axis = (inst->attrs[0]).as_or_throw<ffi::String>();
if (thread_axis != axis) {
return std::nullopt;
}
- return
Downcast<IntImm>(sch->Get(Downcast<LoopRV>(inst->inputs[0]))->extent)->value;
+ return
(sch->Get((inst->inputs[0]).as_or_throw<LoopRV>())->extent).as_or_throw<IntImm>()->value;
Review Comment:
### Readability: Redundant Parentheses
The expression has multiple redundant parentheses around `inst->inputs[0]`
and the `sch->Get(...)` call. Removing them improves readability significantly.
```suggestion
return
sch->Get(inst->inputs[0].as_or_throw<LoopRV>())->extent.as_or_throw<IntImm>()->value;
```
##########
src/s_tir/meta_schedule/mutator/mutate_thread_binding.cc:
##########
@@ -114,9 +114,9 @@ std::vector<MutateThreadBindingNode::Candidate>
MutateThreadBindingNode::FindCan
}
TVM_FFI_ICHECK_EQ(inst->inputs.size(), 1);
TVM_FFI_ICHECK_EQ(inst->attrs.size(), 1);
- if (Downcast<ffi::String>(inst->attrs[0]) != "threadIdx.x") return false;
+ if ((inst->attrs[0]).as_or_throw<ffi::String>() != "threadIdx.x") return
false;
- return
sampled_split_insts.find(Downcast<s_tir::LoopRV>(inst->inputs[0]).get()) !=
+ return
sampled_split_insts.find((inst->inputs[0]).as_or_throw<s_tir::LoopRV>().get())
!=
sampled_split_insts.end();
Review Comment:
### Readability: Redundant Parentheses
The redundant parentheses around `inst->attrs[0]` and `inst->inputs[0]` can
be removed to make the code cleaner.
```suggestion
if (inst->attrs[0].as_or_throw<ffi::String>() != "threadIdx.x") return
false;
return
sampled_split_insts.find(inst->inputs[0].as_or_throw<s_tir::LoopRV>().get()) !=
sampled_split_insts.end();
```
##########
src/relax/transform/inline_functions.cc:
##########
@@ -214,8 +214,10 @@ Pass InlinePrivateFunctions() {
}
auto write_ptr = mod.CopyOnWrite();
- for (const auto& [key, func] : replacements) {
- write_ptr->Remove(Downcast<GlobalVar>(key));
+ for (const auto& [gvar, base_func] : mod->functions) {
+ if (replacements.count(gvar)) {
+ write_ptr->Remove(gvar);
+ }
}
Review Comment:
### Critical Bug: Potential Iterator Invalidation
Iterating over `mod->functions` while mutating the module in-place via
`write_ptr->Remove(gvar)` is extremely dangerous.
If `mod` is not shared (i.e., its `use_count` is 1), `mod.CopyOnWrite()`
will return the same underlying `IRModuleNode` without copying. In this
scenario, `write_ptr->Remove(gvar)` directly erases elements from the same
`functions` map that is currently being iterated over, leading to **iterator
invalidation** and undefined behavior (typically a crash or silent corruption).
To prevent this, you should collect the `GlobalVar`s to be removed into a
temporary `std::vector` first, and then perform the removal in a separate loop.
```c
std::vector<GlobalVar> to_remove;
for (const auto& [gvar, base_func] : mod->functions) {
if (replacements.count(gvar)) {
to_remove.push_back(gvar);
}
}
for (const auto& gvar : to_remove) {
write_ptr->Remove(gvar);
}
```
##########
src/s_tir/meta_schedule/schedule_rule/multi_level_tiling_tensor_core.cc:
##########
@@ -446,11 +446,11 @@ std::vector<State>
MultiLevelTilingTensorCoreNode::TransformIntermediateOutputLa
// Get the shape of the wmma accumulator
auto [frag_shape_m, frag_shape_n] = [&]() {
tirx::SBlock intrin_block =
- Downcast<tirx::SBlockRealize>(
-
tirx::TensorIntrin::Get(state->intrin_group.init_intrin).value()->desc->body)
+
(tirx::TensorIntrin::Get(state->intrin_group.init_intrin).value()->desc->body)
+ .as_or_throw<tirx::SBlockRealize>()
->block;
- tirx::For loop_m = Downcast<tirx::For>(intrin_block->body);
- tirx::For loop_n = Downcast<tirx::For>(loop_m->body);
+ tirx::For loop_m = (intrin_block->body).as_or_throw<tirx::For>();
+ tirx::For loop_n = (loop_m->body).as_or_throw<tirx::For>();
Review Comment:
### Readability: Redundant Parentheses
The redundant parentheses around the `desc->body` access and the block body
accesses can be removed to simplify the code.
```suggestion
tirx::TensorIntrin::Get(state->intrin_group.init_intrin).value()->desc->body
.as_or_throw<tirx::SBlockRealize>()
->block;
tirx::For loop_m = intrin_block->body.as_or_throw<tirx::For>();
tirx::For loop_n = loop_m->body.as_or_throw<tirx::For>();
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
