Lunderberg opened a new pull request, #15577:
URL: https://github.com/apache/tvm/pull/15577
Prior to this commit, a `relax.PrimValue` could have a datatype, but
couldn't have a corresponding `tir.PrimExpr`. As a result, it could not be
used to specify tensor shapes. This makes some expressions require fallback to
`R.Tensor(ndim=ndim)`, even though the shape could still be inferred.
```python
@R.function
def func(
A: R.Tensor(16, 16),
first_n_rows: R.prim("int64"),
) -> R.Tensor([first_n_rows, 16]):
# ^^^^^^^^^^^^
# R.Tensor requires a PrimExpr, not relax.Expr
#
# Operations may require PrimExpr
# vvvvvvvvvvvv
out = R.op.strided_slice(axis=[0], begin=[0], end=[first_n_rows])
return out
```
This commit adds a `Optional<PrimExpr> value` field to the `PrimStructInfo`.
This field acts similarly to the `PrimExpr` fields already used in
`ShapeStructInfo`, and may contain symbolic variables.
```python
@R.function
def func(
A: R.Tensor(16, 16),
# TIR definitions in signature allow in-line definitions,
# similar to R.Tensor and R.Shape. R.Prim takes `dtype` or
# `value` kwarg to distinguish between in-line symbolic variable
# and string representation of dtype.
first_n_rows: R.prim(value="first_n_rows_tir"),
) -> R.Tensor(["first_n_rows_tir", 16]):
# Body contains a TIR variable definition, which may be used
# in function calls, inferred shape annotations.
first_n_rows_tir = T.int64()
out = R.op.strided_slice(axis=[0], begin=[0], end=[first_n_rows])
return out
```
--
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]
