This is an automated email from the ASF dual-hosted git repository.
yongwww pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 64fe5a8a89 [Unity][DistIR] Add DTensor struct info propagation rule
for stop_lift_params (#16170)
64fe5a8a89 is described below
commit 64fe5a8a89696d382c400203e152d3ca55bc565e
Author: Hongyi Jin <[email protected]>
AuthorDate: Tue Nov 28 16:59:04 2023 -0500
[Unity][DistIR] Add DTensor struct info propagation rule for
stop_lift_params (#16170)
* stop lift param dtensor struct info
* format
---
src/relax/op/distributed/op.cc | 10 ++++++++++
1 file changed, 10 insertions(+)
diff --git a/src/relax/op/distributed/op.cc b/src/relax/op/distributed/op.cc
index 2c742595a3..ef780c6df8 100644
--- a/src/relax/op/distributed/op.cc
+++ b/src/relax/op/distributed/op.cc
@@ -37,6 +37,16 @@ StructInfo InferDistStructInfoCallTIR(const Call& call,
const BlockBuilder& ctx)
TVM_REGISTER_OP("relax.call_tir")
.set_attr<FInferStructInfo>("dist.FInferStructInfo",
InferDistStructInfoCallTIR);
+StructInfo InferDistStructInfoStopLiftParams(const Call& call, const
BlockBuilder& ctx) {
+ if (call->args.size() != 1) {
+ ctx->ReportFatal(Diagnostic::Error(call) << "stop_lift_params should have
exact 1 arg.");
+ }
+ return Downcast<StructInfo>(call->args[0]->struct_info_.value());
+}
+
+TVM_REGISTER_OP("relax.builtin.stop_lift_params")
+ .set_attr<FInferStructInfo>("dist.FInferStructInfo",
InferDistStructInfoStopLiftParams);
+
} // namespace distributed
} // namespace relax
} // namespace tvm
