[TIR] add unit-tests for upcoming primfunc-slicing #12794
Conversation
cconvey
force-pushed
the
call-tir-tests1
branch
from
841a32d to
27a14f3
Compare
cconvey
force-pushed
the
call-tir-tests1
branch
3 times, most recently
from
fe15a2d to
2edbc2d
Compare
cconvey
changed the title
cconvey
changed the title
cconvey
force-pushed
the
call-tir-tests1
branch
from
2edbc2d to
a416a4e
Compare
cconvey
changed the title
|
@JosephTheOctonaut @Lunderberg : mind reviewing? |
|
@JosephTheOctonaut @Lunderberg : Feel free to only give cursory reviews on this. It's mostly meant as a rough roadmap. I expect each individual unit test to need a little cleanup as it gets enabled. |
Lunderberg
left a comment
Lunderberg
left a comment
There was a problem hiding this comment.
Looks good overall for mapping out the general features for TIR-to-TIR calls.
| # 3) As support for 'call_tir' becomes more complete, this test should once again | ||
| # fail, because the specified callee doesn't exist. This test should be updated | ||
| # to once again expect failure. | ||
| @pytest.mark.skip(reason="Awaiting TVMScript support for 'call_tir' token.") |
There was a problem hiding this comment.
Would recommend using pytest.mark.xfail instead of pytest.mark.skip. That way the test still executes, but is allowed to fail without failing the test suite as a whole. It also can be useful during development, if overlapping tests start to pass after implementing a new feature, because they are listed as XPASS in the summary.
There was a problem hiding this comment.
I've updated this to be @pytest.mark.xfail(reason=..., strict=True) based on suggestions from TVM CI people (@driazati et al.)
| # | ||
| # NOTE! The role of this unit test should evolve as follows: | ||
| # 1) Initially the test should fail, because we haven't yet changed the TVMScript | ||
| # parser to support 'call_tir'. |
There was a problem hiding this comment.
Possibly a larger question, possibly bikeshedding: Should we change the name to tir_subroutine in order to avoid confusion with Relax's call_tir intrinsic?
There was a problem hiding this comment.
100% agree that the name call_tir should be revisited, but IMHO we should answer a few more questions before picking the name. (E.g., are we sure this is any different from call_external once the callee has been extracted?)
Would it be okay for now to just add a clarifying comment?
| # def main(): | ||
| # A = T.buffer[[1], "int8"] | ||
| # A[0] = 0 | ||
| # with T.block(): |
There was a problem hiding this comment.
Where did the name of this block go?
| # with T.block(): | ||
| # call_tir(add_one, A) | ||
| # | ||
| # # NOTE: it's not entirely clear how the name for the generated callee is chosen |
There was a problem hiding this comment.
We could generate a name from the name of the function (must be unique within the IRModule), and the name of the block (must be unique within each function). There could still be name collisions if the module already contains a function named f"{func_name}_subroutine_{block_name}" (or whatever specific convention we want), but those feel like they'd be rare enough that we could check and append a number afterwards in those cases
There was a problem hiding this comment.
Would it make sense, at least initially, to have people specify the callee subroutine name as part of the TVMScript?
E.g,. T.annotate("extract_as_subroutine", "some_name")?
We could add more defaulting behavior later if/when desired. But this would give us more flexibility during pathfinding.
There was a problem hiding this comment.
That sounds reasonable to me, that it could be specified explicitly.
| # def add_one(X: T.buffer[[1], "int8"]): | ||
| # X[0] += 1 | ||
| # | ||
| # def create_schedule(self, irmod: tvm.ir.module.IRModule, blockname_to_funcname_map:Map[str,str]) -> tvm.tir.Schedule: |
There was a problem hiding this comment.
What is the create_schedule function intended to do?
There was a problem hiding this comment.
I was (perhaps mistakenly) thinking that that slicing would involve the application of scheduling primitives. I've removed this pseudocode in the next commit, since that's a premature assumption.
| reason="Awaiting TVMScript support for lowering of 'T.call_tir' to 'T.call_packed'." | ||
| ) | ||
| class TestLowerCallTir(tvm.testing.CompareBeforeAfter): | ||
| # @tvm.script.ir_module |
There was a problem hiding this comment.
Will also need a transform defined here. I think we'll want it to occur in tvm.tir.transform.MakePackedAPI.
There was a problem hiding this comment.
Thanks, I've added a code comment about this.
| # A = T.buffer[[1], "int8"] | ||
| # A[0] = 0 | ||
| # with T.block(): | ||
| # # TODO: figure out the right TVMScript thing to do here |
There was a problem hiding this comment.
We can use the function calls currently generated by SplitHostDevice as a template (link). Overall, we'll want to output a Call node with the operation builtin::tvm_call_packed().
There was a problem hiding this comment.
Thanks, I've added a code comment about this.
|
Note to self: consider renaming "call_tir" to something else, e.g. "call_tir_subroutine", to avoid confusing overlap with Relax's "call_tir" token. |
cconvey
force-pushed
the
call-tir-tests1
branch
2 times, most recently
from
33f3c3a to
3c066c2
Compare
Lunderberg
left a comment
Lunderberg
left a comment
There was a problem hiding this comment.
Thank you for making the changes and documenting which parts may change as the design involves. LGTM!
Add unit tests (initially disabled) to motivate upcoming work on semi-automated slicing of primfuncs. (I.e., extracting some subtree of a primfunc body's TIR into a separate primfunc.)
cconvey
force-pushed
the
call-tir-tests1
branch
from
3c066c2 to
0ab0b99
Compare
[TIR] Add disabled primfunc-slice unit tests Add unit tests (initially disabled) to motivate upcoming work on semi-automated slicing of primfuncs. (I.e., extracting some subtree of a primfunc body's TIR into a separate primfunc.)
2 participants
Introduces (disabled) unit tests that can gradually be enabled as we make progress on primfunc slicing (a.k.a. "functionalization").
cc @Hzfengsy @junrushao1994