commit | fddbec7079a817d3339b45fedd3c5b8326cfafac | [log] [tgz] |
---|---|---|
author | Krzysztof Parzyszek <kparzysz@quicinc.com> | Tue Jul 11 11:37:51 2023 -0500 |
committer | GitHub <noreply@github.com> | Tue Jul 11 11:37:51 2023 -0500 |
tree | d513b419fb5b26ea11721ccc8cf0cded2cc0929f | |
parent | 592b3583dc403af2ffc9601b50339c1f84ce2227 [diff] |
[TIR] Implement TIR macros (#15260) * [TIR] Implement TIR macros 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. * Fix linter * Fix linter again One linter suggested something that the other didn't like... * Get rid of T.insert, apply macro via function-call syntax * Store closure vars in TIRMacro * ast.parse always returns ast.Module, hence doc is doc.Module * Simplify `expand_macro`, capture environment variables * Implement macro hygiene * Fix linter * Make T.macro work same as T.macro() The previous commit inadvertently made T.macro (without parentheses) illegal, only abbreviated form allowed was T.macro(). Restore T.macro as a valid decorator use. * Edit comment: insertion -> expansion * Add import pytest * One more typo... * Remove stale testcase
Documentation | Contributors | Community | Release Notes
Apache TVM is a compiler stack for deep learning systems. It is designed to close the gap between the productivity-focused deep learning frameworks, and the performance- and efficiency-focused hardware backends. TVM works with deep learning frameworks to provide end to end compilation to different backends.
TVM is licensed under the Apache-2.0 license.
Check out the TVM Documentation site for installation instructions, tutorials, examples, and more. The Getting Started with TVM tutorial is a great place to start.
TVM adopts apache committer model, we aim to create an open source project that is maintained and owned by the community. Check out the Contributor Guide.
We learned a lot from the following projects when building TVM.