Hzfengsy opened a new pull request #10307:
URL: https://github.com/apache/tvm/pull/10307
This PR tries to simplify `compute_at` loop nesting with a static bound. For
example:
```Python
@T.prim_func
def static_bound(A: T.Buffer[(32, 1), "float32"], C: T.Buffer[(32, 1),
"float32"]) -> None:
B = T.alloc_buffer((32, 1), "float32")
for i, j in T.grid(32, 1):
with T.block("B"):
vi = T.axis.spatial(32, i)
vj = T.axis.spatial(1, j)
B[vi, vj] = A[vi, vj] * 2.0
for i, j in T.grid(32, 32):
with T.block("C"):
vi = T.axis.spatial(32, i)
vj = T.axis.spatial(1, j)
T.where(j < 1)
C[vi, vj] = B[vi, vj] + 1.0
```
if we try to compute_at `block B` under loop `i`, we will get
```Python
@T.prim_func
def func(A: T.Buffer[(32, 1), "float32"], C: T.Buffer[(32, 1), "float32"])
-> None:
B = T.alloc_buffer([32, 1], dtype="float32")
for i in T.serial(32):
for ax0, ax1 in T.grid(1, 32):
with T.block("B"):
vi = T.axis.spatial(32, i + ax0)
vj = T.axis.spatial(1, ax1)
T.where(ax1 < 1)
B[vi, vj] = A[vi, vj] * T.float32(2)
for j in T.serial(32):
with T.block("C"):
vi = T.axis.spatial(32, i)
vj = T.axis.spatial(1, j)
T.where(j < 1)
C[vi, vj] = B[vi, vj] + T.float32(1)
```
However, we can make loop `ax1` be a loop with extent 1 since the bound `ax1
< 1` is static (Please see the testcase)
cc @spectrometerHBH @wrongtest @junrushao1994
--
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]
