gemini-code-assist[bot] commented on code in PR #20009:
URL: https://github.com/apache/tvm/pull/20009#discussion_r3587968129
##########
src/relax/distributed/transform/lower_global_view_to_local_view.cc:
##########
@@ -403,8 +403,14 @@ class LowerTIRToLocalView : public ExprMutator {
ffi::Array<Expr> args = val->args[1].as_or_throw<Tuple>()->fields;
for (const auto& arg : args) {
const auto* ty = GetTypeAs<DTensorTypeNode>(arg);
- TVM_FFI_ICHECK(ty);
- sharding_specs.push_back(ShardingSpec(ty->device_mesh, ty->placement));
+ if (ty) {
+ sharding_specs.push_back(ShardingSpec(ty->device_mesh, ty->placement));
+ } else {
+ TVM_FFI_ICHECK(GetType(arg).as<PrimTypeNode>())
+ << "Expected call_tir arguments to be distributed tensors or
primitive arguments, "
+ "but got "
+ << arg << " with type " << GetType(arg);
Review Comment:
The well-formed check here only validates that non-distributed arguments
have `PrimTypeNode`. However, in Relax, shape expressions (represented as
`ShapeExpr` with `ShapeTypeNode`) are also valid primitive/shape arguments in
`call_tir`. If a `ShapeExpr` is passed, this check will fail and crash the
compiler. We should allow both `PrimTypeNode` and `ShapeTypeNode` to ensure
robustness.
```suggestion
TVM_FFI_ICHECK(GetType(arg).as<PrimTypeNode>() ||
GetType(arg).as<ShapeTypeNode>())
<< "Expected call_tir arguments to be distributed tensors or
primitive arguments, "
"but got "
<< arg << " with type " << GetType(arg);
```
##########
src/relax/transform/fuse_tir.cc:
##########
@@ -833,17 +804,19 @@ class FusedTIRConstructor : public ExprVisitor {
void MapInputBuffer(const tirx::PrimFunc& func, const relax::Expr& args) {
ffi::Array<Expr> arg_list;
ffi::Array<tirx::Buffer> buffer_list;
- if (const auto* arg_tuple = args.as<TupleNode>()) {
- arg_list = arg_tuple->fields;
- } else {
- arg_list = {args};
- }
+ ffi::Array<Expr> call_args = args.as_or_throw<Tuple>()->fields;
- TVM_FFI_ICHECK_GE(func->params.size(), arg_list.size());
- for (size_t i = 0; i < arg_list.size(); ++i) {
+ TVM_FFI_ICHECK_GE(func->params.size(), call_args.size());
+ for (size_t i = 0; i < call_args.size(); ++i) {
+ const Expr& arg = call_args[i];
const tirx::Var& param = func->params[i];
- const tirx::Buffer& buffer = func->buffer_map.at(param);
- buffer_list.push_back(buffer);
+ if (GetType(arg).as<PrimTypeNode>()) {
+ func_info_.symbolic_var_matcher.Match(param.as_or_throw<PrimExpr>(),
+ arg.as_or_throw<PrimExpr>());
+ } else {
+ arg_list.push_back(arg);
+ buffer_list.push_back(func->buffer_map.at(param));
+ }
Review Comment:
Checking `GetType(arg).as<PrimTypeNode>()` to distinguish scalar/symbolic
arguments from buffer arguments is fragile. Scalar/shape arguments can also be
represented as `ShapeExpr` (which has `ShapeTypeNode` instead of
`PrimTypeNode`). If a `ShapeExpr` is passed, it will fall into the `else`
branch, causing `func->buffer_map.at(param)` to throw an out-of-range exception
and crash the compiler.
A much more robust and idiomatic way to distinguish buffer parameters from
scalar parameters is to check if the parameter is present in the PrimFunc's
`buffer_map` (i.e., `func->buffer_map.count(param)`). All buffer parameters of
a PrimFunc must be mapped in `buffer_map`, whereas scalar parameters are not.
```suggestion
if (func->buffer_map.count(param)) {
arg_list.push_back(arg);
buffer_list.push_back(func->buffer_map.at(param));
} else {
func_info_.symbolic_var_matcher.Match(param.as_or_throw<PrimExpr>(),
arg.as_or_throw<PrimExpr>());
}
```
--
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]
