cconvey commented on code in PR #12794: URL: https://github.com/apache/tvm/pull/12794#discussion_r981496165
########## tests/python/unittest/test_slice_tir.py: ########## @@ -0,0 +1,178 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import tvm +import tvm.testing +from tvm.script import tir as T +import pytest + + +# ABOUT THIS FILE: +# We (cconvey / OctoML) are working on a sequence of PRs to allow a single TIR primfunc's +# AST to be sliced into multiple partitiones, where each partition will be converted into +# a new TIR primfunc. (See https://en.wikipedia.org/wiki/Program_slicing). +# +# The unit tests below provide a roadmap for that sequence of PRs; each PR should allow +# one more of these tests to pass. +# +# NOTE: These unit tests may change as work progresses. They aren't meant to +# indicate hard requirements. + +# NOTE! The `tvm.testing.CompareBeforeAfter` class provides TWO useful mechanisms for +# these tests: +# +# (a) It lets us specify code snippets which are valid Python, but which aren't YET +# recognized as valid TVMScript. This allows unit tests for new constructs, +# e.g. 'call_tir(...)' to simply be disabled rather than fully commented out. +# +# (b) It lets us structurally compare the TIR bodies of two primfuncs. +# +# Note that some of the tests below will require the structural comparison of +# two entire IRModules, not just primfuncs. This will require adding functionality +# to the `CompareBeforeAfter` class, or implementing that level of comparison within +# the individual unit tests. +# +# Some of the unit tests below which require whole-IRModule comparison. For expedience +# we simply comment out the (early draft) bodies of those unit tests, rather than +# hacking their structure to get the benefits of (a). + + +# step 1: that vvvv simple passes Python / TVMScript parsing. +# +# The only requirement for this test is that the TVMScript parser +# doesn't raise an error when encountering `T.call_tir(foo)`, +# where "foo" is a syntactically valid TVMScript function name. +# +# 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'. +# +# 2) Initial TVMScript support for 'call_tir' will be minimal, essentially ignoring +# it. This test should pass once that change is made. +# +# 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. [email protected](reason="Awaiting TVMScript support for 'call_tir' token.") +class TestParseCallTIR(tvm.testing.CompareBeforeAfter): + """ + Simply confirm that the TIR node `call_tir` doesn't interfere with + the successful parsing of the TVMScript. + """ + + def before(): + T.call_tir(add_one) + T.evalute(0) + + def expected(): + T.evaluate(0) + + +# Step 2: transform annotated block ==> separate primfuncs + call_tir [email protected]( + reason="Awaiting TVMScript support for 'call_tir' and T.annotation(\"functionalize\")." +) +class TestAnnotateAndSliceTIR(tvm.testing.CompareBeforeAfter): + # def test_annotate_and_slice(): + # @tvm.script.ir_module + # class irmod_before: + # @T.prim_func + # def main(A: T.Buffer[(1,), "int8"): + # #A = T.match_buffer(a, (1,), "int8") + # A[0] = 0 + # with T.block("block_foo"): # optional: give this block a name, perhaps for testing? + # # NOTE: nice to have: human control over name used for the generated callee + # T.annotate("functionalize") # TODO: find the actually correct name (and/or update the name in Eric's commit) + # A[0] += 1 + # return 42 + # + # @tvm.script.ir_module + # class irmod_after: + # @T.prim_func + # def main(): + # A = T.buffer[[1], "int8"] + # A[0] = 0 + # with T.block(): + # call_tir(add_one, A) + # + # # NOTE: it's not entirely clear how the name for the generated callee is chosen Review 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. -- 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]
