tqchen commented on code in PR #20009:
URL: https://github.com/apache/tvm/pull/20009#discussion_r3588051177
##########
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:
Thanks for flagging this. I verified that a grouped ShapeExpr is not a valid
positional argument in the migrated ABI: one ShapeExpr field may represent
multiple TIR scalar parameters, while each call_tir tuple field maps to exactly
one leading PrimFunc parameter. The well-formed checker rejects that case, and
FuseTIR already unpacks valid Relax shape parameters into individual PrimExpr
fields before constructing call_tir. I updated distributed lowering to classify
each position from the PrimFunc buffer_map and to require an individual
PrimExpr for scalar parameters, avoiding a crash or silent positional shift. I
also added a focused well-formedness regression for grouped ShapeExpr arguments.
##########
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:
Thanks, I adopted the buffer_map-based classification. FuseTIR now
determines tensor versus scalar handling from the corresponding PrimFunc
parameter instead of the Relax argument type, and emits a clear TypeError if a
scalar parameter receives anything other than an individual PrimExpr. This
preserves the one-field/one-parameter ABI; valid Relax shape parameters
continue to be unpacked before call_tir construction. The focused FuseTIR and
well-formedness regressions pass.
--
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]
