Lunderberg commented on code in PR #12794: URL: https://github.com/apache/tvm/pull/12794#discussion_r975504368
########## 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.") Review 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. ########## 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(): Review Comment: We probably should add a comment here that this implementation will require `tvm.testing.CompareBeforeAfter` to support comparisons between `IRModule` as well as `PrimFunc`. ########## 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: 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 ########## 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. + """ Review Comment: The `CompareBeforeAfter` utility needs a `transform` to be defined. In this case, it would be simplest to define an identity transformation (i.e `transform = tvm.tir.transform.prim_func_pass(lambda func, _mod,_ctx: func, 0)`). ########## 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 + # @T.prim_func + # 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: + # sch = irmod.create_schedule() + # for blockname, new_funcname in blockname_to_funcname_map.items(): + # sch.annotate(blockname, new_funcname, True) + # # TODO: + pass + + +# Step 3: transform call_tir ==> packed call [email protected]( + reason="Awaiting TVMScript support for lowering of 'T.call_tir' to 'T.call_packed'." +) +class TestLowerCallTir(tvm.testing.CompareBeforeAfter): + # @tvm.script.ir_module Review Comment: Will also need a `transform` defined here. I think we'll want it to occur in `tvm.tir.transform.MakePackedAPI`. ########## 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(): Review Comment: Where did the name of this block go? ########## 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'. Review 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? ########## 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 + # @T.prim_func + # 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: + # sch = irmod.create_schedule() + # for blockname, new_funcname in blockname_to_funcname_map.items(): + # sch.annotate(blockname, new_funcname, True) + # # TODO: + pass + + +# Step 3: transform call_tir ==> packed call [email protected]( + reason="Awaiting TVMScript support for lowering of 'T.call_tir' to 'T.call_packed'." +) +class TestLowerCallTir(tvm.testing.CompareBeforeAfter): + # @tvm.script.ir_module + # class test_lower_before: + # @T.prim_func + # def main(): + # A = T.buffer[[1], "int8"] + # A[0] = 0 + # with T.block(): + # call_tir(add_one, A) + # + # @T.prim_func + # def add_one(X: T.buffer[[1], "int8"]): + # X[0] += 1 + # + # @tvm.script.ir_module + # class test_lower_after: + # @T.prim_func + # def main(): + # A = T.buffer[[1], "int8"] + # A[0] = 0 + # with T.block(): + # # TODO: figure out the right TVMScript thing to do here Review Comment: We can use the function calls currently generated by `SplitHostDevice` as a template ([link](https://github.com/apache/tvm/blob/main/src/tir/transforms/split_host_device.cc#L336)). Overall, we'll want to output a Call node with the operation `builtin::tvm_call_packed()`. ########## 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) Review Comment: The name we used was `"extract_as_subroutine"`, but it's not one that we're heavily tied to if we want to change it. ########## 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 + # @T.prim_func + # 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: Review Comment: What is the `create_schedule` function intended to do? -- 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]
