[TIR][UX] Implement privacy annotations in TIR #15214

slyubomirsky · 2023-07-03T18:38:09Z

Implements proposal #14899 for TIR. Adapted from Unity PR #15171, by the same method as PR #15140. Namely, a TIR PrimFunc can be specified to be private (without a global symbol attribute) in TVMScript in the prim_func decorator. By default, PrimFuncs are not private, so they will have a global_symbol attribute that is mapped to their name.

Example usage:

# not private: its global symbol will be "func"
@T.prim_func
def func(...): ...

# no global symbol included
@T.prim_func(private=True)
def func(...): ...

This did require changing very, very many tests, unfortunately.

tvm-bot · 2023-07-03T18:38:12Z

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

cc @Hzfengsy, @junrushao, @quic-sanirudh, @shingjan _{See #10317 for details}

_{Generated by tvm-bot}

slyubomirsky · 2023-07-07T00:34:08Z

@junrushao Any insight on the test failures involving check_sketches (e.g., in tests/python/unittest/test_meta_schedule_schedule_rule_cross_thread_reduction.py)? Getting rid of the global symbols in those tests didn't affect them, so I'm not sure where my changes could affect these

slyubomirsky · 2023-07-14T22:56:25Z

The test errors turned out to global symbols not matching within the meta scheduler (thanks @vinx13 for the help). Hopefully will be solved shortly.

slyubomirsky · 2023-07-18T16:36:37Z

@tqchen @vinx13 @junrushao Now that the tests pass, I would appreciate close review, as there were very many tests changed and there seemed to be a lot of functionality affected by global symbols.

include/tvm/script/ir_builder/tir/ir.h

junrushao · 2023-07-18T19:21:01Z

src/ir/function.cc

      LOG(FATAL) << "Do not support function type " << func->GetTypeKey();
    });

+TVM_REGISTER_GLOBAL("ir.BaseFuncWithoutAttr")


Nice! Looks like we always forgot to upstream this piece from the unity branch

python/tvm/tir/schedule/testing.py

…rm_make_unpacked_api

slyubomirsky · 2023-07-19T15:51:17Z

@tvm-bot rerun

junrushao

This is extremely meaningful contribution to overall clarity and formalization of TIR! LGTM

Fix FuseOps to adapt #15137 Fix TIR TVMScript to adapt #15214

Required for internal functions after PR apache#15214

Analogous to #14901, treat GlobalVar callees as internal function calls in CodeGenC. This specific PR doesn't provide new end-to-end functionality, as the target="c" backend isn't compiled. It does lead into allowing subroutines in any target whose codegen derives from CodeGenC, which will depend on the single-module lowering flow in #14985. * [CodeGenC] Added unit tests for desired behavior * [CodeGenC] Handle GlobalVar callee as internal function call * Update CodeGenC subclasses for updated interface - Call `DeclareFunction` for each `PrimFunc`, prior to any `AddFunction` calls - Provide both `GlobalVar` and `PrimFunc` to `AddFunction` calls. * Updated CRT test to expect forward declaration * Provide forward declarations for call_extern in cmsis * Avoid duplicate forward declaration C's automatic pointer cast (e.g. `void*` to `int*`) means that use of the arguments to infer the function signature may be incorrect. If a `call_extern` refers to a function within the same module, only output a single forward declaration based on the PrimFunc's parameters, not based on the CallNode's arguments. * Updated expected ptx cuda * Cast the AOT pools to the arg type * Improved tvm::GetType for tvm_access_ptr and address_of These `Call` instances can return a `PointerType(PrimType(pointee_dtype))` rather than a `PrimType(DataType::Handle())`. * [ARM][Topi] Update micro kernels to use same argument type as caller Previously, the micro kernels for gemm, avg_pool, max_pool, and tensordot relied on C's implicit type conversions for the arguments, when the caller's argument types differ from the signature's parameter types. This works, except when the codegen has auto-generated a forward declaration based on the caller's argument types, such as during AOT, which then causes a conflicting definition. Since the codegen cannot determine the functions names from the `"pragma_import_c"` in order to suppress these forward declarations, this conflict can be more easily resolved by updating the micro kernel signatures. The three types of mismatches are below. - Use of `int` or `long` parameters, whose width may vary by compiler, instead of fixed-width types. - TIR expecting the data array's integer type to also be used as an error code's return type, rather than the micro kernels' `int32_t` error code. - Pointer conversion done during argument conversion. Type conversions are done at the start of each micro kernel, to avoid changing types that are used within the computational sections of each micro kernel. * Updated unit tests with private=True Required for internal functions after PR #15214 * Docstring updates from review

slyubomirsky mentioned this pull request Jul 3, 2023

[Unity][UX][TIR] Implement privacy annotation for the @prim_func decorator #15171

Closed

slyubomirsky requested review from junrushao and tqchen July 3, 2023 18:44

slyubomirsky force-pushed the privacy-annotation-tir-mainline branch from 487615b to 72a542e Compare July 14, 2023 19:02

junrushao reviewed Jul 18, 2023

View reviewed changes

include/tvm/script/ir_builder/tir/ir.h Outdated Show resolved Hide resolved

junrushao reviewed Jul 18, 2023

View reviewed changes

python/tvm/tir/schedule/testing.py Outdated Show resolved Hide resolved

slyubomirsky added 19 commits July 18, 2023 20:41

Implement privacy annotation in TIR

cc2bb1f

Fix broken TIR tests

bf1d06f

Treat it as an error to specify the global symbol if marked private

a17ae54

Handle global symbol in testing utils

119ced2

Transfer global symbol in SplitHostDevice

caa19a8

Preserve global symbol in USMP convert pool allocations

b19cd9d

Use global symbol in numerous tests

182de83

Remove some code duplication

80b4418

formatting

a7e86f7

pylint

1543dce

Formatting

eedb556

Undo SplitHostDevice changes, use privacy annotation in tests instead

be9ee73

Fix rebase mistakes in test_tir_schedule_pad_einsum

98bd20f

Factor out func_name in function.cc

9bf0c12

Fix global symbol in more tests

a7ca19a

Fix more global symbols in unit tests

cc7f39c

Fix global symbols in test_pass_plan_devices.py::test_lowered

c07a149

Fix global symbol in test_tir_transform_helpers

57bbceb

Fix test_tir_schedule_cache_read_write tests

9777abf

slyubomirsky added 11 commits July 18, 2023 20:42

Fix global symobls in test_tvmscript_parser_tir.py

3b9246b

Transfer over without_attr from Unity branch

8b219e5

ignore global symbols in meta schedule tests

4a2e3b2

Do not use global symols in test_tir_transform_make_unpacked_pi.py

6ee9ce9

Fix global symbol in test_evaluator_with_preproc

f2f3976

Linting fix

3ab5307

Fix unused import

745421b

Do not use private with an explicit global symbol in test_tir_transfo…

4baeacd

…rm_make_unpacked_api

Use an ordinary bool for the PrimFuncFrame's is_private attribute

64c0cf0

Do not use abbreviation in assert_structural_equal_gs (expand name)

2f70c39

Fix test case in test_tir_transform_split_host_device

645d356

slyubomirsky force-pushed the privacy-annotation-tir-mainline branch from acfb358 to 645d356 Compare July 19, 2023 02:06

junrushao approved these changes Jul 20, 2023

View reviewed changes

junrushao merged commit 7e830e5 into apache:main Jul 20, 2023

MasterJH5574 added a commit to mlc-ai/relax that referenced this pull request Aug 1, 2023

Fix TIR TVMScript to adapt apache/tvm#15214

e1e5333

tqchen pushed a commit that referenced this pull request Aug 1, 2023

[MERGE-FIX] Update the code to fix merge issues

53fd712

Fix FuseOps to adapt #15137 Fix TIR TVMScript to adapt #15214

Lunderberg added a commit to Lunderberg/tvm that referenced this pull request Aug 3, 2023

Updated unit tests with private=True

e77fa25

Required for internal functions after PR apache#15214

ysh329 mentioned this pull request Oct 18, 2023

[Release] v0.14.0 Release Candidate Notes #15948

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[TIR][UX] Implement privacy annotations in TIR #15214

[TIR][UX] Implement privacy annotations in TIR #15214

Uh oh!

slyubomirsky commented Jul 3, 2023

Uh oh!

tvm-bot commented Jul 3, 2023

Uh oh!

slyubomirsky commented Jul 7, 2023

Uh oh!

slyubomirsky commented Jul 14, 2023

Uh oh!

slyubomirsky commented Jul 18, 2023

Uh oh!

Uh oh!

junrushao Jul 18, 2023

Uh oh!

slyubomirsky Jul 18, 2023

Uh oh!

Uh oh!

slyubomirsky commented Jul 19, 2023

Uh oh!

junrushao left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

[TIR][UX] Implement privacy annotations in TIR #15214

[TIR][UX] Implement privacy annotations in TIR #15214

Uh oh!

Conversation

slyubomirsky commented Jul 3, 2023

Uh oh!

tvm-bot commented Jul 3, 2023

Uh oh!

slyubomirsky commented Jul 7, 2023

Uh oh!

slyubomirsky commented Jul 14, 2023

Uh oh!

slyubomirsky commented Jul 18, 2023

Uh oh!

Uh oh!

junrushao Jul 18, 2023

Choose a reason for hiding this comment

Uh oh!

slyubomirsky Jul 18, 2023

Choose a reason for hiding this comment

Uh oh!

Uh oh!

slyubomirsky commented Jul 19, 2023

Uh oh!

junrushao left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants