Hzfengsy opened a new pull request #8354:
URL: https://github.com/apache/tvm/pull/8354
This PR enables meta_programming for TVM Functions (usually used for
TVMScript), which is useful to define only one function(TVMScript) but used in
different shapes.
Example:
Mate fucntion
```Python
@tvm.script.tir
def mem_copy(a: ty.handle, b: ty.handle, m: ty.int32, n: ty.int32) -> None:
A = tir.match_buffer(a, (m, n), "float32")
B = tir.match_buffer(b, (m, n), "float32")
with tir.block([m, n], "") as [vi, vj]:
B[vi, vj] = A[vi, vj]
```
Instruction:
```Python
a, _, m, n = mem_copy.params
func = mem_copy.specialize({a: tir.decl_buffer((16, 16))})
# or
func = mem_copy.specialize({n: 16, m: 16})
```
Specialized function:
```Python
@tvm.script.tir
def mem_copy_16_16(a: ty.handle, b: ty.handle) -> None:
A = tir.match_buffer(a, (16, 16), "float32")
B = tir.match_buffer(b, (16, 16), "float32")
with tir.block([16, 16], "") as [vi, vj]:
B[vi, vj] = A[vi, vj]
```
cc @tqchen @junrushao1994 @MasterJH5574 @comaniac
--
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]
