masahi commented on code in PR #14548:
URL: https://github.com/apache/tvm/pull/14548#discussion_r1162242916
##########
python/tvm/relax/transform/legalize_ops/index.py:
##########
@@ -59,3 +60,39 @@ def _strided_slice(bb: BlockBuilder, call: Call) -> Expr:
call.attrs.axes,
slice_mode="end",
)
+
+
+@register_legalize("relax.dynamic_strided_slice")
+def _dynamic_strided_slice(bb: BlockBuilder, call: Call) -> Expr:
+ # 1. Insert shape function
+ output_shape = bb.normalize(
+ bb.call_te(
+ topi.shape_func_dynamic_strided_slice,
Review Comment:
> If this is true, I don't think this is currently supported in TVMScript
hmm interesting. Don't we create TVMScript with shapes like (`n`, 16) all
the time? I hope the output of a shape func can also be expressed like that.
> Also, if we go with TVMScript, do we require users to write the TIR
primfunc by hand?
Yes but writing in TVMScript is only required if a user wants to exploit
"partially-static" shape. I think this is a relatively rare and advanced use
case. `te.compute`-based definition should always be possible.
And other than dynamic slicing, I cannot think of other ops that may want to
implement their shape func in TVMScript. Maybe `reshape` and other
shape-changing ops might apply but I haven't met a real-world use case.
--
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]
