yzh119 opened a new pull request, #14316:
URL: https://github.com/apache/tvm/pull/14316
# Motivation
Currently, we miss a schedule primitive to change the data type of allocated
buffer (e.g. via `cache_read`/`cache_write`), and thus we cannot perform type
conversion while loading data from global to shared memory.
This PR adds a new schedule primitive `set_dtype` that follows the interface
of `set_scope` and allows users to customize the allocated buffers' data type.
# Example
Before running `set_dtype`:
```python
@T.prim_func
def main(A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128),
"float32")):
# with T.block("root"):
B = T.alloc_buffer((128, 128), "float16")
for i, j in T.grid(128, 128):
with T.block("B"):
vi, vj = T.axis.remap("SS", [i, j])
T.reads(A[vi, vj])
T.writes(B[vi, vj])
B_subregion0 = T.match_buffer(B[vi, vj], (), "float16",
offset_factor=1)
B_subregion0[()] = T.Cast("float16", A[vi, vj] * T.float32(2))
for i, j in T.grid(128, 128):
with T.block("C"):
vi, vj = T.axis.remap("SS", [i, j])
T.reads(B[vi, vj])
T.writes(C[vi, vj])
B_subregion1 = T.match_buffer(B[vi, vj], (), "float16",
offset_factor=1)
C[vi, vj] = T.Cast("float32", B_subregion1[()]) + T.float32(1)
```
then we perform the `set_dtype` schedule:
```python
sch = tir.Schedule(before_set_dtype)
sch.set_dtype("B", buffer_index=0, dtype="float16")
print(sch.mod["main"].script())
```
we get transformed code:
```python
@T.prim_func
def after_set_dtype(
A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
) -> None:
B = T.alloc_buffer((128, 128), dtype="float16")
for i, j in T.grid(128, 128):
with T.block("B"):
vi, vj = T.axis.remap("SS", [i, j])
B[vi, vj] = T.cast(A[vi, vj] * 2.0, "float16")
for i, j in T.grid(128, 128):
with T.block("C"):
vi, vj = T.axis.remap("SS", [i, j]
C[vi, vj] = T.cast(B[vi, vj], "float32") + 1.0
```
where data type conversions are inserted automatically.
# Other Usage
Using the combination of `cache_read` + `set_dtype` can help us load data
from the memory hierarchy while converting data to the desired type.
cc @Hzfengsy @vinx13 @junrushao
--
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]
