kparzysz-quic opened a new pull request, #15260:
URL: https://github.com/apache/tvm/pull/15260
This patch introduces two new symbols: `T.macro` and `T.insert`. `T.macro`
is a decorator that, when applied to a function, turns the body of that
function into a piece of TIR that can be inserted via `T.insert` into a
PrimFunc.
For example:
```python
@T.macro
def copy_backwards(dst, src, size):
with T.block("backwards"):
for i in T.serial(size):
ai = T.axis.remap("S", [i])
T.reads(src[0:size])
T.writes(dst[0:size])
dst[ai] = src[size - ai - 1]
@T.prim_func
def foo_int32(A: T.Buffer((128,), "int32"), B: T.Buffer((128,), "int32")):
T.insert(copy_backwards, A, B, 128)
@T.prim_func
def foo_int8(A: T.Buffer((128,), "int8"), B: T.Buffer((128,), "int8")):
T.insert(copy_backwards, A, B, 128)
```
The above will generate two PrimFuncs that do the same backwards copy, but
applied to buffers with different data types.
Semantics:
- Function that is decorated with @T.macro can have any parameters that
follow Python syntax, i.e. positional, keyword, etc. Type annotations are not
required, but are allowed.
- The arguments to `T.insert` are macro name followed by the argument list.
For `T.insert(arg1, arg2, arg3, ...)`, the values are substituted into the body
of the macro as in the call `arg1(arg2, arg3, ...)`. The body with the
substituted values is then inserted at the point where the `T.insert` is
located.
--
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]
