Lunderberg commented on code in PR #16642:
URL: https://github.com/apache/tvm/pull/16642#discussion_r1535744877
##########
src/tir/ir/function.cc:
##########
@@ -21,12 +21,52 @@
* \file src/tir/ir/function.cc
* \brief The function data structure.
*/
+#include <tvm/relax/struct_info.h>
#include <tvm/runtime/registry.h>
+#include <tvm/tir/analysis.h>
#include <tvm/tir/function.h>
#include <tvm/tir/op.h>
namespace tvm {
namespace tir {
+namespace {
+relax::StructInfo InferStructInfo(const PrimFunc& prim_func) {
+ Array<relax::StructInfo> params;
+ for (const auto& param : prim_func->params) {
+ relax::StructInfo param_sinfo = [&]() -> relax::StructInfo {
+ if (auto opt_buf = prim_func->buffer_map.Get(param)) {
+ auto buf = opt_buf.value();
+ relax::ShapeExpr shape(
+ buf->shape.Map([](PrimExpr dim) { return cast(DataType::Int(64),
dim); }));
+ return relax::TensorStructInfo(shape, buf->dtype);
+ }
+
+ if (auto prim_type = param->type_annotation.as<PrimTypeNode>();
Review Comment:
That's a good call, and I've updated it to provide `R.Object` instead of
`R.Tensor` in this case. My main goal with this conditional was to avoid
having `DLTensor*` parameters be labeled with `R.Prim(dtype='handle')`, as that
would be incompatible with `R.Tensor` arguments.
In the future, the annotations could be improved by inspecting the PrimFunc
body. If the handle-typed `tir.Var` is used in a context that requires a
`DLTensor*` argument (typically `builtin::tvm_struct_get`), then the annotation
could be improved from `R.Object` to `R.Tensor`. That would allow better
type-checking at the Relax callsite, but isn't necessary for a first
implementation.
--
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]
