This PR adds boolean operators to OperationDoc. This is needed by the TIR expression printing because it has `tir::And` and `tir::Or`.
Tracking issue: #11912
This PR:
- Adds an entry point for the TVMScript Unified Printer
- Adds a helper object class `RootNodeContainer` to provide an injection point for the actual printer implementation to add specialized logic on the root node to print.
Tracking issue: https://github.com/apache/tvm/issues/11912
* [TVMScript] IRBuilder, IRBuilderFrame base class
This PR introduces basic data structures of the generic IRBuilder
across the codebase.
IRBuilder is a general-purpose IRBuilder that can be used in TIR, Relax
and any other vendor-specific dialects; IRBuilderFrame is where contexual
information as stored in the IRBuilder.
* fix linter
* Update include/tvm/script/ir_builder/base.h
Co-authored-by: Junru Shao <junrushao1994@gmail.com>
This PR:
- Adds VarTable for the new TVMScript Printer
Compared to the prototype version, this:
- Removes unnecessary public methods.
- GetObjectName
- GetUniqueName
- Add Frame parameter for `Define` methods. VarTable will add callback to Frame to remove variable when Frame exits.
- Changes DocFactory from `ExprDoc(ObjectPath)` to `ExprDoc()` to simplify var definition.
Tracking issue: https://github.com/apache/tvm/issues/11912
This PR:
- Implement Frame for the TVMScript Unified Printer
Compared to the prototype version, this:
- Removes the dependency of VarTable (SymbolTable) from Frame
- Adds a callback array to the Frame base class so that VarTable can add callback to clean variable when Frame goes out scope
Tracking issue: https://github.com/apache/tvm/issues/11912
This adds an ability to print a "diagnostic marker" based on a given ObjectPath. For example, say we are printing a fragment of TIR like
```
for i in T.serial(10):
a[i] = 5
```
and we would like bring the user's attention to the bound of the loop:
```
for i in T.serial(10):
^^
a[i] = 5
```
In this case we would give the doc printer an object path that represents this loop bound, i.e. something like `path_to_underline=ObjectPath.root().attr("extent")`
Tracking issue: https://github.com/apache/tvm/issues/11912
Storage rewrite was missing a visitor for let statements so buffers
added in them would still refer to the pre-rewritten version. This error
was originally noticed when using `global.vtcm` buffers which get
changed to let statements by LowerVtcmAlloc.
Implementing the test for this change also required adding support for
vectorized datatypes to tvmscript. The solution included is a little
hacky and involes adding the datatypes to the `global()` table of each
module they need to be defined in.
This PR:
- Add the source_paths attribute to Doc base class.
- Add the corresponding Python binding for it.
This PR is depended by multiple tasks, including the diagnostic output in DocPrinter, VarTable and IRDocisifer.
Tracking issue: https://github.com/apache/tvm/issues/11912
Co-authored-by: Greg Bonik <gbonik@octoml.ai>
- test_aot_legalize_packed_call.py: `T.preflattened_buffer` returns `void`
- test_tir_intrin.py: `type` here should be `buffer_type`
- test_tir_transform_flatten_buffer.py: `extents` should be `list`
- test_tir_transform_hoist_expression.py: change `tir` into `T` and register `Let` expression in `script/tir/intrin.py`
- test_tir_transform_storage_flatten.py: `T.allocate` has no argument named `strides`
Previously we use ANSI color sequences to highlight TVM script. In jupyter notebook environments, such color sequence will be recoginized and translated to corresponding HTML to display things.
This works fine for most notebook environments (including the jupyter notebook and the VS Code plugin). Recently, thanks to @tqchen, we found that Google Colab does not well support ansi color sequence for 24-bit colors (`JupyterLight` and `VSCDark`) that all its displayed colors are unexpectedly black/gray/white. To also bring highlighting in Colab, in this PR, we directly render the highlighted code with HTML when a notebook environment is detected.
* feat(ux): highlight tvm script
* resolve dependency
* refact(tvmscript): highlight fallback as plain text with warning; put ansi colors as default terminal style
* refact(tvmscript): make terminal style close to the default notebook style
* fix(pylint): disable=import-outside-toplevel
* fix(ci-dep): Pygments>=2.4.0 to support ansicolors w/o #
* refact: making Pygments versioning most robust and user-friendly
* fix: pylint var naming
This PR addes:
- All ExprDoc subclasses
- Their Python bindings
- Support of ExprDoc in PythonDocPrinter
- Unit tests for ExprDoc in PythonDocPrinter
Tracking issue: https://github.com/apache/tvm/issues/11912
* Compute common type for shape elements in BroadcastHelper
The corresponding dimensions in the input/output tensors in a broadcast
operations may have the same value, but different types (e.g. int32 vs
int64).
When the broadcast helper tries to unify the dimensions it also needs
to compute the common type to hold the dimension.
* Cast and simplify both members of `Range`
Only the `min` member was type-casted, which could lead to ranges with
different types for `min` and `extent`.
Move the casts to the argument of Simplify, so that they can be eliminated
if they aren't needed.
* Type-check iv domain ranges, use cast only if needed in MakeLoopNest
In some cases the domain ranges had the `min` and the `extent` values
be of different types (e.g. [(int64)0, 32)). This is an error, and it
can lead to compilation failures later on. Add a check for equal types
here to catch this early.
Also, only add the cast operation when the desired type differs from
the current one to keep the expressions simpler.
* Check that variable and substituted expression have same types
Add a check to IRSubstitute to detect when the type of a variable and
the type of the expression to replace it with have different types.
* Add testcase
* [TVMScript] Use void for lambda parameters, allow mismatch in Substitute
When the script parser deals with lambdas, it creates Var objects for each
parameter. Their actual types are not known at the time, and the properly
typed variables are subtituted in the body later. Since the default dtype
of a Var is "int32", this could lead to a type mismatch in Substitute.
To deal with this scenario, use "void" for newly created Vars in the
parser, and add an exception to Substitute to allow replacing void Vars
with expressions of any type.
* Fix type error in test_reduce_combiner_simplify
* Restart CI
Co-authored-by: Jiawei Liu <jaway.liu@gmail.com>
* [TIR.Constant] U1 usecase
Constants are now aggregated into one struct and initialized in default_lib0.c
file
Change-Id: I34d61f8139c8a92c06944fe990ba892a660476fd
Unit test fixed
Change-Id: I436e7b6d6b3064b3f8bbfbb048d4296b63a6b69c
* Refactored
Addressed:
* PoolInfo splitted to WorkspacePoolInfo and ConstantPoolInfo
* workspace_byte_alignment moved to ExecutorCodegenMetadata
* getModuleAlignment -> GetModuleAlignment
* GenerateInternalWorkspaceBuffers refactored
* reverted format change of src/tir/transforms/legalize_packed_calls.cc
* addressed comments for src/tir/usmp/analysis/extract_buffer_info.cc
* removed commented code from include/tvm/tir/usmp/utils.h
Change-Id: I7d1b32884b0e5992e2e00c7838c85e425d9c25fd
* more unit test fixes
Change-Id: I573a05fa1cb4037ae83691f7dff2c2724b1d7700
* More refactoring and unit test fixes
Added ConstantMemoryPools
Change-Id: If1e391c631575980564bca790ba33748c82d907f
* bugfix
Change-Id: Iacc7a9d734a505dfa0d8d32d23ea3f57e6de8582
* refactoring. added constant_alignment
added constant_alignment
unit tests updated
Change-Id: I378193cb9e675e352c61d96ff4e09655090053e1
* unit-test bugix
Change-Id: Ia4411d59c4a376c01326fed366cdb196a432899e
* unit test fix
Change-Id: Ia2077bdeb1d2c6c9827eeef90ab410ae31b8c4a4
* Added support for c++ runtime
* refactored
* renamed pools and consts
renamed pools and consts to workspace_pools and constant_pools
* addressed upstream comments
* addressed upstream comments-2
* addressed upstream comments-3
* [TVMSCRIPT] Improve tvmscript type hints
- Change numeric types to classes so they work as function arguments.
- Add var as a class.
- Add floordiv, index, and mod to PrimExpr.
* use Union
* [TVMScript] Allow T.Buffer[] arg annotation to use int as shape
Both the function `tvm.tir.decl_buffer` and the TVMScript
`T.match_buffer` expression allow a `PrimExpr` to be passed as the buffer
shape, which is interpreted as a 1-d buffer of that size. This allows
the same behavior to be used in the `T.Buffer` syntactic sugar.
(e.g. `A: T.Buffer[16, "float32"]` instead of `A: T.Buffer[(16,), "float32"`)
* Fixed round-trip when buffer size contains an expression
* [TVMScript] Support function call to help construct AST
* add test
* update test
* more comment
* fix for avoiding Buffer.vload(...) case
* update parse error msg
* wrap func call with try / catch, emit error msg
* silence pylint
* [TVMScript] Allow `val = buf[index]` without type annotation
Other instances of `var = expr` were previously allowed without
requiring a type annotation, by using the dtype of the expression as
the dtype of `var`. This behavior didn't work for `buf[index]`
expressions, which are internally represented as `BufferSlice` python
objects, and only converted to `BufferLoad` primexprs when used as an
expression.
This commit adds a `dtype` property to `BufferSlice`, allowing
`buf[index]` to be used in a let statement without a type annotation.
* Reverted a wider change
Automatically adding a type annotation to Var if it could be
determined from the dtype let the unit test directly compare the
annotated and unannotated versions of buffer load. Unfortunately, it
also broke 54 unrelated tests, so that change is removed from this PR.
* add get_c_struct_name() method to Metadata to distinguish struct type name in llvm
* add metadata serialization support to llvm codegen
* Organize MetadataQueuer into a separate file.
* Add DiscoverArraysVisitor to metadata_utils
* Fill DLTensor metadata in LegalizePackedCalls.
* Improve error message from Call asserts
* Pass non-String device_context down to codegen.
* this is necessary to allow CodeGenCPU to emit calls that include resource_handle.
* Scope usage of lvalue refs in LowerTVMBuiltin to avoid corrupt memory.
* test fixes
* Also fill preflattened_buffer_map (TODO, maybe don't do this)
* Fix C codegen.
* Set USMP elem_offset to 0.
* Clarify calculation of byte_offset from elem_offset.
* fix tests
* Fix arm compile warning
* Fix hexagon test.
* previously I believe we required interface_api == "c", but
this really means to generate C API bindings, and we are generating
"packed" bindings.
* I think "c" was chosen here because the distinction between
interface-api and use-unpacked-api is confusing. "c" interface-api
means to generate an entrypoint API for microcontrollers that
accepts bare data buffers. "packed" interface-api means to generate
a TVMBackendPackedCFunc entrypoint. use-unpacked-api forms the same
determination for the operator functions.
* A further confusion here is that there are two ways to call
"packed" operator functions: tir.tvm_builtin_call_packed and
tir.tvm_builtin_call_cpacked. This distinction describes whether or
not to late-bind calls via TVMBackendGetFuncFromEnv. Right now, AOT
only ever requires call_cpacked because target_host == target, and
for all suitable target_host, we expect a single DSO-exportable
runtime.Module. When we move away from this by introducing
heterogeneous target support to AOT, we can use this as a condition
to help us choose between call_cpacked and call_packed (and
possibly add a compile-time option to assert it is call_cpacked,
for situations where we really don't want call_packed).
* Document T.preflattened_buffer
* Fix test_aot_legalize_packed_calls
* Address manupa comments
* Fix convert_pool_allocations_to_offsets test.
* lint
* Fix T.preflattened_buffer
* Add preflattened_buffer_map to TIRTextPrinter
* Fix tests
* Fix BYOC
* Fix invoking C device API.
* remove comments
* Address Mousius comments
* lint
* lint
* Fix GMock linking on new CMake
* address masahi comment
Co-authored-by: Masahiro Masuda <masahi129@gmail.com>
* Respect dtype in Scalarize.
* Add unittest.
* Fix lint.
* Promote dtype of IntImm to match loop_var in For.
* Fix dtype mismatches.
* Lint
* Lint.
* jostle ci
* Match dtype in hybrid parser.
This demonstrates how to selectively extract and tune tasks from a whole relay mod, and apply the tuned schedule during the final `relay.build(...)`.
This flow is entirely different from existing tests in `test_meta_schedule_tune_relay.py` where ALL ops are extracted and auto-scheduled by MS. My test extracts only int8 `dense` op, applies a manual TIR schedule on it, and leaves int8 `batch_matmul` to be scheduled by TE.
This also serves as an example of autotvm style manual template + tensorization. The manual TIR schedule is equivalent to TE VNNI `dense` schedule in https://github.com/apache/tvm/blob/ce335c3a74185df6cc1152e53c60695d8a418d8e/python/tvm/topi/x86/dense.py#L366-L375
## Context
When dealing with end-to-end models, we note that some tensors may have large shapes. Thus, when designing graph-level IR, we sometimes use `int64` instead of `int32` for the shape. Below is an dense GeMM example which has `int64` input tensor shape:
```python
@tvm.script.ir_module
class Module:
@T.prim_func
def main(rxplaceholder: T.Buffer[(1, 512), "float32"], rxplaceholder_1: T.Buffer[(T.int64(1000), T.int64(512)), "float32"], T_matmul_NT: T.Buffer[(1, T.int64(1000)), "float32"]) -> None:
# function attr dict
T.func_attr({"global_symbol": "dense", "tir.noalias": True, "op_pattern": 3})
# body
# with T.block("root")
for i0_0, i1_0, i0_1, i1_1, i2_0, i0_2, i1_2, i2_1, i0_3, i1_3 in T.grid(1, 4, 1, 25, 8, 1, 10, 64, 1, 1):
with T.block("T_matmul_NT"):
i = T.axis.spatial(1, 0)
j = T.axis.spatial(T.int64(1000), i1_0 * T.int64(250) + i1_1 * T.int64(10) + i1_2)
k = T.axis.reduce(512, i2_0 * 64 + i2_1)
T.reads(T_matmul_NT[i, j], rxplaceholder[i, k], rxplaceholder_1[j, k])
T.writes(T_matmul_NT[i, j])
T.block_attr({"layout_free_placeholders":[rxplaceholder_1], "meta_schedule.tiling_structure":"SSRSRS"})
with T.init():
T_matmul_NT[i, j] = T.float32(0)
T_matmul_NT[i, j] = T_matmul_NT[i, j] + rxplaceholder[i, k] * rxplaceholder_1[j, k]
```
## Problem
Though our TVMScript printer can easily print `int64` constants, the parser had poor support for `int64`. So this PR introduces some parser support for `int64`, basically about the data type of loop variables, block iterators and block read/write regions.
Besides the parser, most of the TIR schedule primitives didn't take `int64` into account in their implementations. These schedule primitives will be fixed and updated in recent future, in followup PRs.