psrivas2 opened a new pull request, #14066:
URL: https://github.com/apache/tvm/pull/14066
This change adds a PrimFunc level analysis to propose layout transformations
to block and buffers in the PrimFunc based on the layout transformations to
PrimFunc outputs. It analyzes buffer access to figure this out. It tries to
preserve sequential access to buffers when it does this.
For example given the following PrimFunc and write buffer "relu"
transformation `lambda n, c, h, w: (n, h, w, c // 4, c % 4)`, it will suggest
to make the following transformations.
* Block transformation on "compute": `lambda n, c, h, w: (n, h, w, c // 4, c
% 4)`
* Buffer transformation on "relu": `lambda n, c, h, w: (n, h, w, c // 4, c
% 4)`
* Buffer transformation on "arg": `lambda n, c, h, w: (n, h, w, c // 4, c %
4)`
```python
@T.prim_func
def elemwise_relu(
arg: T.Buffer((32, 64, 224, 224), "float32"),
relu: T.Buffer((32, 64, 224, 224), "float32"),
):
for i0, i1, i2, i3 in T.grid(32, 64, 224, 224):
with T.block("compute"):
v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
T.reads(arg[v_i0, v_i1, v_i2, v_i3])
T.writes(relu[v_i0, v_i1, v_i2, v_i3])
relu[v_i0, v_i1, v_i2, v_i3] = T.max(arg[v_i0, v_i1, v_i2,
v_i3], T.float32(0))
```
These transformations can then be applied on the PrimFunc to get the
PrimFunc with new layout.
```python
@T.prim_func
def elemwise_relu(
arg: T.Buffer((32, 224, 224, 16, 4), "float32"),
relu: T.Buffer((32, 224, 224, 16, 4), "float32"),
):
for ax0, ax1, ax2, ax3, ax4 in T.grid(32, 224, 224, 16, 4):
with T.block("compute"):
v0, v1, v2, v3, v4 = T.axis.remap("SSSSS", [ax0, ax1, ax2, ax3,
ax4])
T.reads(arg[v0, v1, v2, v3, v4])
T.writes(relu[v0, v1, v2, v3, v4])
relu[v0, v1, v2, v3, v4] = T.max(arg[v0, v1, v2, v3, v4],
T.float32(0))
```
--
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]
