You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "slyubomirsky (via GitHub)" <gi...@apache.org> on 2023/12/05 04:12:27 UTC
[PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
slyubomirsky opened a new pull request, #16204:
URL: https://github.com/apache/tvm/pull/16204
This PR implements a pass that automatically extracts DataflowBlocks out of binding blocks (i.e., groups consecutive pure operations in binding blocks into a dataflow block).
Note that this pass makes use of `CanonicalizeBindings` and in the process of testing it, I encountered bugs in `CanonicalizeBindings` that I deal with in here:
1. `CanonicalizeBindings` dealt with the case of a var that appeared only inside a DataflowBlock (demoted to a dataflow var), but it made a mistake in cases where a var was used only inside DataflowBlocks, plural, meaning outside the one where it was defined (thus it should be exposed as an output and thus not demoted to a dataflow var).
2. `CanonicalizeBindings` did not handle inner functions correctly. It will correctly recurse into inner functions and also treat any captured vars as appearing outside their original dataflow block.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416177769
##########
tests/python/relax/test_transform_extract_dataflow_blocks.py:
##########
@@ -0,0 +1,268 @@
+# 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 import relax
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+
+class ExtractCompare(tvm.testing.CompareBeforeAfter):
+ transform = relax.transform.ExtractDataflowBlocks()
+
+
+# functions that will not change
+class TestTrivial(ExtractCompare):
+ @I.ir_module
+ class Before:
+ # already a DF block
+ @R.function
+ def main(A: R.Tensor, B: R.Tensor):
+ with R.dataflow():
+ x = R.add(A, B)
+ y = R.multiply(x, A)
+ z = R.add(x, y)
+ q = R.multiply(y, z)
+ p = R.add(z, q)
+ R.output(p)
+ return p
+
+ # too small
+ @R.function
+ def func(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ x = R.add(A, B)
+ y = R.subtract(x, B)
+ return y
+
+ # too few pure ops between non-dataflow ops
+ @R.function(pure=False)
+ def func2(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ _ = R.print(format="Hi there!")
+ y = R.add(A, B)
+ _ = R.print(y, format="Sum: {}")
+ x = R.multiply(y, y)
+ if R.const(False):
+ _ = R.print(format="True branch")
+ q = R.add(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ else:
+ _ = R.print(format="False branch")
+ q = R.subtract(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ p = R.multiply(w, w)
+ return p
+
+ Expected = Before
+
+
+class TestBasic(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ return v
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ return v
+
+
+class TestMultipleBlocks(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ _ = R.print(format="Hi mom!")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ return d
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ _ = R.print(format="Hi mom!")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ R.output(d)
+ return d
+
+
+class TestExtractInsideBranches(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ if R.const(True):
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ else:
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, b)
+ d = R.add(c, c)
+ e = R.multiply(d, d)
+ return e
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+
+ if R.const(True):
+ with R.dataflow():
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ R.output(b)
+ # weird but the parser requires this construct
+ c = b
+ else:
+ with R.dataflow():
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ R.output(b)
+ c = b
+ with R.dataflow():
+ d = R.multiply(c, c)
+ e = R.add(d, d)
+ f = R.multiply(e, e)
+ R.output(f)
+ return f
+
+
+class TestTreatNonCallAsPure(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(t: R.Tuple(R.Tensor, R.Tensor)) -> R.Tensor:
+ x = t[0]
+ y = t[1]
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ return w
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(t: R.Tuple(R.Tensor, R.Tensor)) -> R.Tensor:
+ with R.dataflow():
+ x = t[0]
+ y = t[1]
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ R.output(w)
+ return w
+
+
+class TestInnerFunction(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
Review Comment:
You're right, the pass does care about purity, so I'll add those cases.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416176948
##########
tests/python/relax/test_transform_extract_dataflow_blocks.py:
##########
@@ -0,0 +1,268 @@
+# 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 import relax
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+
+class ExtractCompare(tvm.testing.CompareBeforeAfter):
+ transform = relax.transform.ExtractDataflowBlocks()
+
+
+# functions that will not change
+class TestTrivial(ExtractCompare):
+ @I.ir_module
+ class Before:
+ # already a DF block
+ @R.function
+ def main(A: R.Tensor, B: R.Tensor):
+ with R.dataflow():
+ x = R.add(A, B)
+ y = R.multiply(x, A)
+ z = R.add(x, y)
+ q = R.multiply(y, z)
+ p = R.add(z, q)
+ R.output(p)
+ return p
+
+ # too small
+ @R.function
+ def func(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ x = R.add(A, B)
+ y = R.subtract(x, B)
+ return y
+
+ # too few pure ops between non-dataflow ops
+ @R.function(pure=False)
+ def func2(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ _ = R.print(format="Hi there!")
+ y = R.add(A, B)
+ _ = R.print(y, format="Sum: {}")
+ x = R.multiply(y, y)
+ if R.const(False):
+ _ = R.print(format="True branch")
+ q = R.add(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ else:
+ _ = R.print(format="False branch")
+ q = R.subtract(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ p = R.multiply(w, w)
+ return p
+
+ Expected = Before
+
+
+class TestBasic(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ return v
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ return v
+
+
+class TestMultipleBlocks(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ _ = R.print(format="Hi mom!")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ return d
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ _ = R.print(format="Hi mom!")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ R.output(d)
+ return d
+
+
+class TestExtractInsideBranches(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ if R.const(True):
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ else:
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, b)
+ d = R.add(c, c)
+ e = R.multiply(d, d)
+ return e
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+
+ if R.const(True):
+ with R.dataflow():
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ R.output(b)
+ # weird but the parser requires this construct
+ c = b
+ else:
+ with R.dataflow():
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ R.output(b)
+ c = b
+ with R.dataflow():
+ d = R.multiply(c, c)
+ e = R.add(d, d)
+ f = R.multiply(e, e)
+ R.output(f)
+ return f
+
+
+class TestTreatNonCallAsPure(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(t: R.Tuple(R.Tensor, R.Tensor)) -> R.Tensor:
+ x = t[0]
Review Comment:
Yeah, good idea.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on PR #16204:
URL: https://github.com/apache/tvm/pull/16204#issuecomment-1841610319
Thanks for the excellent feedback, @Lunderberg. I will also add a few more test cases per your suggestions.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416620121
##########
tests/python/relax/test_transform_convert_dataflow.py:
##########
@@ -0,0 +1,493 @@
+# 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 import relax
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+
+class ExtractCompare(tvm.testing.CompareBeforeAfter):
+ transform = relax.transform.ConvertToDataflow()
+
+
+# functions that will not change
+class TestTrivial(ExtractCompare):
+ @I.ir_module
+ class Before:
+ # already a DF block
+ @R.function
+ def main(A: R.Tensor, B: R.Tensor):
+ with R.dataflow():
+ x = R.add(A, B)
+ y = R.multiply(x, A)
+ z = R.add(x, y)
+ q = R.multiply(y, z)
+ p = R.add(z, q)
+ R.output(p)
+ return p
+
+ # too small
+ @R.function
+ def func(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ x = R.add(A, B)
+ return x
+
+ # too few pure ops between non-dataflow ops
+ @R.function(pure=False)
+ def func2(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ _ = R.print(format="Hi there!")
+ y = R.add(A, B)
+ _ = R.print(y, format="Sum: {}")
+ x = R.multiply(y, y)
+ if R.const(False):
+ _ = R.print(format="True branch")
+ q = R.add(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ else:
+ _ = R.print(format="False branch")
+ q = R.subtract(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ p = R.multiply(w, w)
+ return p
+
+ Expected = Before
+
+
+class TestBasic(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ return v
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ return v
+
+
+class TestMultipleBlocks(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ _ = R.print(format="Hi mom!")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ return d
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ _ = R.print(format="Hi mom!")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ R.output(d)
+ return d
+
+
+class TestExtractInsideBranches(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ if R.const(True):
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ else:
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, b)
+ d = R.add(c, c)
+ e = R.multiply(d, d)
+ return e
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+
+ if R.const(True):
+ with R.dataflow():
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ R.output(b)
+ # weird but the parser requires this construct
+ c = b
+ else:
+ with R.dataflow():
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ R.output(b)
+ c = b
+ with R.dataflow():
+ d = R.multiply(c, c)
+ e = R.add(d, d)
+ f = R.multiply(e, e)
+ R.output(f)
+ return f
+
+
+class TestTreatNonCallAsPure(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def tuples_and_const(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ t1 = (x, y, x)
+ t2 = (y, y, x)
+ c = R.const([1, 2, 3], dtype="int32")
+ return c
+
+ @R.function
+ def shapes() -> R.Shape:
+ s1 = R.shape((1, 2, 3))
+ s2 = R.shape((4, 5, 6))
+ s3 = R.shape((7, 8, 9))
+ return s3
+
+ @R.function
+ def prim_values():
Review Comment:
Whoa, that's an impressive segfault. It looks like the sequence of events is as follows:
1. We hit [this TypeError](https://github.com/apache/tvm/blob/unity/python/tvm/script/parser/relax/entry.py#L452), which gives a diagnostic that the `R.Prim` needs to have one of the two argument types.
2. The TypeError is caught [here](https://github.com/apache/tvm/blob/unity/python/tvm/script/parser/relax/parser.py#L102), where it is passed to the diagnostic context to pretty-print and re-throw.
3. After showing the error message, [this conditional](https://github.com/apache/tvm/blob/unity/src/ir/diagnostic.cc#L129) resets the diagnostic renderer and raises an exception.
4. The DiagnosticError is caught [here](https://github.com/apache/tvm/blob/unity/python/tvm/script/parser/relax/parser.py#L111), where it is passed to the diagnostic context to pretty-print and re-throw.
5. When reaching [here](https://github.com/apache/tvm/blob/unity/src/ir/diagnostic.cc#L118), `(*this)->renderer` has been replaced with a default constructed instance. The default-constructed instance from step (3) has a nullptr for the renderer, which causes the segfault.
It looks like there's a lot of `except DiagnosticError: raise` scattered throughout the parser to prevent this from occurring (by preventing step (4) from occurring), but these `eval_struct_info` and `eval_struct_info_proxy` don't have it. That seems like a rather fragile way to avoid some pretty opaque errors, so I'm wondering if we want to change that in the long term.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416184540
##########
tests/python/relax/test_transform_extract_dataflow_blocks.py:
##########
@@ -0,0 +1,268 @@
+# 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 import relax
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+
+class ExtractCompare(tvm.testing.CompareBeforeAfter):
+ transform = relax.transform.ExtractDataflowBlocks()
+
+
+# functions that will not change
+class TestTrivial(ExtractCompare):
+ @I.ir_module
+ class Before:
+ # already a DF block
+ @R.function
+ def main(A: R.Tensor, B: R.Tensor):
+ with R.dataflow():
+ x = R.add(A, B)
+ y = R.multiply(x, A)
+ z = R.add(x, y)
+ q = R.multiply(y, z)
+ p = R.add(z, q)
+ R.output(p)
+ return p
+
+ # too small
+ @R.function
+ def func(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ x = R.add(A, B)
+ y = R.subtract(x, B)
+ return y
+
+ # too few pure ops between non-dataflow ops
+ @R.function(pure=False)
+ def func2(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ _ = R.print(format="Hi there!")
+ y = R.add(A, B)
+ _ = R.print(y, format="Sum: {}")
+ x = R.multiply(y, y)
+ if R.const(False):
+ _ = R.print(format="True branch")
+ q = R.add(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ else:
+ _ = R.print(format="False branch")
+ q = R.subtract(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ p = R.multiply(w, w)
+ return p
+
+ Expected = Before
+
+
+class TestBasic(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ return v
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ return v
+
+
+class TestMultipleBlocks(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ _ = R.print(format="Hi mom!")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ return d
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ _ = R.print(format="Hi mom!")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ R.output(d)
+ return d
+
+
+class TestExtractInsideBranches(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ if R.const(True):
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ else:
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, b)
+ d = R.add(c, c)
+ e = R.multiply(d, d)
+ return e
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+
+ if R.const(True):
+ with R.dataflow():
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ R.output(b)
+ # weird but the parser requires this construct
+ c = b
+ else:
+ with R.dataflow():
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ R.output(b)
+ c = b
+ with R.dataflow():
+ d = R.multiply(c, c)
+ e = R.add(d, d)
+ f = R.multiply(e, e)
+ R.output(f)
+ return f
+
+
+class TestTreatNonCallAsPure(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(t: R.Tuple(R.Tensor, R.Tensor)) -> R.Tensor:
+ x = t[0]
+ y = t[1]
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ return w
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(t: R.Tuple(R.Tensor, R.Tensor)) -> R.Tensor:
+ with R.dataflow():
+ x = t[0]
+ y = t[1]
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ R.output(w)
+ return w
+
+
+class TestInnerFunction(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ @R.function(pure=False)
+ def inner_func(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(x, z)
+ v = R.add(y, w)
+ _ = R.print(format="oops")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.multiply(a, b)
+ return c
+
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ v = R.divide(w, z)
+ q = inner_func(w, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, a)
+ return c
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+
+ @R.function(pure=False)
+ def inner_func(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(x, z)
+ v = R.add(y, w)
+ R.output(v)
+ _ = R.print(format="oops")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.multiply(a, b)
+ R.output(c)
+ return c
+
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ v = R.divide(w, z)
+ R.output(inner_func, v, w)
+ q = inner_func(w, v)
+ with R.dataflow():
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, a)
+ R.output(c)
+ return c
+
Review Comment:
Great idea. I'll have to handle that case though.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416182676
##########
src/relax/transform/extract_dataflow_blocks.cc:
##########
@@ -0,0 +1,140 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/relax/transform/extract_dataflow_blocks.cc
+ * \brief Pass for extracting groups of pure operations without
+ * dataflow into dataflow blocks.
+ */
+
+#include <tvm/relax/expr.h>
+#include <tvm/relax/expr_functor.h>
+#include <tvm/relax/transform.h>
+#include <tvm/relax/utils.h>
+
+namespace tvm {
+namespace relax {
+
+class DataflowBlockExtractor : public ExprMutator {
+ public:
+ explicit DataflowBlockExtractor(size_t min_size) : ExprMutator(), min_size_(min_size) {}
+
+ Expr VisitExpr_(const SeqExprNode* seq) override {
+ Array<BindingBlock> new_blocks;
+ Expr new_body = VisitExpr(seq->body);
+ bool changed = !new_body.same_as(seq->body);
+ for (auto block : seq->blocks) {
+ BindingBlock new_block = this->VisitBindingBlock(block);
+ changed = changed || !new_block.same_as(block);
+ if (new_block.as<DataflowBlock>()) {
Review Comment:
Good point. I rely on normalization throughout this pass.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416182224
##########
python/tvm/relax/transform/transform.py:
##########
@@ -262,6 +262,23 @@ def LambdaLift() -> tvm.ir.transform.Pass:
return _ffi_api.LambdaLift()
+def ExtractDataflowBlocks(min_size: int = 3) -> tvm.ir.transform.Pass:
Review Comment:
It was an arbitrary call on my part, so sure I can change it to 2.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416181388
##########
src/relax/transform/canonicalize_bindings.cc:
##########
@@ -170,17 +191,26 @@ class CanonicalizePlanner : public ExprVisitor {
}
void VisitExpr_(const VarNode* var) override {
- if (!inside_dataflow_) {
- used_outside_dataflow_.insert(GetRef<Var>(var));
+ auto var_ref = GetRef<Var>(var);
+ // if a var is used in a dataflow block but *not* the one
+ // where it was defined, it also needs to be exposed, so also we treat that as
+ // used outside of a dataflow block
+ if (!inside_dataflow_ ||
+ (def_blocks_.count(var_ref) && !current_block_.same_as(def_blocks_.at(var_ref)))) {
+ used_outside_home_dataflow_.insert(GetRef<Var>(var));
}
}
bool inside_dataflow_{false};
+ BindingBlock current_block_;
Review Comment:
Yeah, I was going to do that for a second reason too: There is no "current block" when you're looking at the `body` field of a SeqExpr. This could lead to weird results.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416390972
##########
tests/python/relax/test_transform_convert_dataflow.py:
##########
@@ -0,0 +1,493 @@
+# 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 import relax
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+
+class ExtractCompare(tvm.testing.CompareBeforeAfter):
+ transform = relax.transform.ConvertToDataflow()
+
+
+# functions that will not change
+class TestTrivial(ExtractCompare):
+ @I.ir_module
+ class Before:
+ # already a DF block
+ @R.function
+ def main(A: R.Tensor, B: R.Tensor):
+ with R.dataflow():
+ x = R.add(A, B)
+ y = R.multiply(x, A)
+ z = R.add(x, y)
+ q = R.multiply(y, z)
+ p = R.add(z, q)
+ R.output(p)
+ return p
+
+ # too small
+ @R.function
+ def func(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ x = R.add(A, B)
+ return x
+
+ # too few pure ops between non-dataflow ops
+ @R.function(pure=False)
+ def func2(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ _ = R.print(format="Hi there!")
+ y = R.add(A, B)
+ _ = R.print(y, format="Sum: {}")
+ x = R.multiply(y, y)
+ if R.const(False):
+ _ = R.print(format="True branch")
+ q = R.add(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ else:
+ _ = R.print(format="False branch")
+ q = R.subtract(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ p = R.multiply(w, w)
+ return p
+
+ Expected = Before
+
+
+class TestBasic(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ return v
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ return v
+
+
+class TestMultipleBlocks(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ _ = R.print(format="Hi mom!")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ return d
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ _ = R.print(format="Hi mom!")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ R.output(d)
+ return d
+
+
+class TestExtractInsideBranches(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ if R.const(True):
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ else:
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, b)
+ d = R.add(c, c)
+ e = R.multiply(d, d)
+ return e
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+
+ if R.const(True):
+ with R.dataflow():
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ R.output(b)
+ # weird but the parser requires this construct
+ c = b
+ else:
+ with R.dataflow():
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ R.output(b)
+ c = b
+ with R.dataflow():
+ d = R.multiply(c, c)
+ e = R.add(d, d)
+ f = R.multiply(e, e)
+ R.output(f)
+ return f
+
+
+class TestTreatNonCallAsPure(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def tuples_and_const(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ t1 = (x, y, x)
+ t2 = (y, y, x)
+ c = R.const([1, 2, 3], dtype="int32")
+ return c
+
+ @R.function
+ def shapes() -> R.Shape:
+ s1 = R.shape((1, 2, 3))
+ s2 = R.shape((4, 5, 6))
+ s3 = R.shape((7, 8, 9))
+ return s3
+
+ @R.function
+ def prim_values():
Review Comment:
It appears that it is not permitted to leave both optional fields empty. I guess it's not such a serious bug if it's not permitted anyway but it would be better to have a nicer error than a segfault.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416389017
##########
tests/python/relax/test_transform_convert_dataflow.py:
##########
@@ -0,0 +1,493 @@
+# 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 import relax
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+
+class ExtractCompare(tvm.testing.CompareBeforeAfter):
+ transform = relax.transform.ConvertToDataflow()
+
+
+# functions that will not change
+class TestTrivial(ExtractCompare):
+ @I.ir_module
+ class Before:
+ # already a DF block
+ @R.function
+ def main(A: R.Tensor, B: R.Tensor):
+ with R.dataflow():
+ x = R.add(A, B)
+ y = R.multiply(x, A)
+ z = R.add(x, y)
+ q = R.multiply(y, z)
+ p = R.add(z, q)
+ R.output(p)
+ return p
+
+ # too small
+ @R.function
+ def func(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ x = R.add(A, B)
+ return x
+
+ # too few pure ops between non-dataflow ops
+ @R.function(pure=False)
+ def func2(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ _ = R.print(format="Hi there!")
+ y = R.add(A, B)
+ _ = R.print(y, format="Sum: {}")
+ x = R.multiply(y, y)
+ if R.const(False):
+ _ = R.print(format="True branch")
+ q = R.add(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ else:
+ _ = R.print(format="False branch")
+ q = R.subtract(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ p = R.multiply(w, w)
+ return p
+
+ Expected = Before
+
+
+class TestBasic(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ return v
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ return v
+
+
+class TestMultipleBlocks(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ _ = R.print(format="Hi mom!")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ return d
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ _ = R.print(format="Hi mom!")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ R.output(d)
+ return d
+
+
+class TestExtractInsideBranches(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ if R.const(True):
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ else:
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, b)
+ d = R.add(c, c)
+ e = R.multiply(d, d)
+ return e
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+
+ if R.const(True):
+ with R.dataflow():
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ R.output(b)
+ # weird but the parser requires this construct
+ c = b
+ else:
+ with R.dataflow():
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ R.output(b)
+ c = b
+ with R.dataflow():
+ d = R.multiply(c, c)
+ e = R.add(d, d)
+ f = R.multiply(e, e)
+ R.output(f)
+ return f
+
+
+class TestTreatNonCallAsPure(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def tuples_and_const(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ t1 = (x, y, x)
+ t2 = (y, y, x)
+ c = R.const([1, 2, 3], dtype="int32")
+ return c
+
+ @R.function
+ def shapes() -> R.Shape:
+ s1 = R.shape((1, 2, 3))
+ s2 = R.shape((4, 5, 6))
+ s3 = R.shape((7, 8, 9))
+ return s3
+
+ @R.function
+ def prim_values():
Review Comment:
Here's a bug I encountered, probably with the parser: Putting `-> R.Prim` as the return would segfault. Maybe it expects `-> R.Prim()`? Not sure why it would segfault though.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416184540
##########
tests/python/relax/test_transform_extract_dataflow_blocks.py:
##########
@@ -0,0 +1,268 @@
+# 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 import relax
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+
+class ExtractCompare(tvm.testing.CompareBeforeAfter):
+ transform = relax.transform.ExtractDataflowBlocks()
+
+
+# functions that will not change
+class TestTrivial(ExtractCompare):
+ @I.ir_module
+ class Before:
+ # already a DF block
+ @R.function
+ def main(A: R.Tensor, B: R.Tensor):
+ with R.dataflow():
+ x = R.add(A, B)
+ y = R.multiply(x, A)
+ z = R.add(x, y)
+ q = R.multiply(y, z)
+ p = R.add(z, q)
+ R.output(p)
+ return p
+
+ # too small
+ @R.function
+ def func(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ x = R.add(A, B)
+ y = R.subtract(x, B)
+ return y
+
+ # too few pure ops between non-dataflow ops
+ @R.function(pure=False)
+ def func2(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ _ = R.print(format="Hi there!")
+ y = R.add(A, B)
+ _ = R.print(y, format="Sum: {}")
+ x = R.multiply(y, y)
+ if R.const(False):
+ _ = R.print(format="True branch")
+ q = R.add(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ else:
+ _ = R.print(format="False branch")
+ q = R.subtract(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ p = R.multiply(w, w)
+ return p
+
+ Expected = Before
+
+
+class TestBasic(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ return v
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ return v
+
+
+class TestMultipleBlocks(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ _ = R.print(format="Hi mom!")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ return d
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ _ = R.print(format="Hi mom!")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ R.output(d)
+ return d
+
+
+class TestExtractInsideBranches(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ if R.const(True):
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ else:
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, b)
+ d = R.add(c, c)
+ e = R.multiply(d, d)
+ return e
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+
+ if R.const(True):
+ with R.dataflow():
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ R.output(b)
+ # weird but the parser requires this construct
+ c = b
+ else:
+ with R.dataflow():
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ R.output(b)
+ c = b
+ with R.dataflow():
+ d = R.multiply(c, c)
+ e = R.add(d, d)
+ f = R.multiply(e, e)
+ R.output(f)
+ return f
+
+
+class TestTreatNonCallAsPure(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(t: R.Tuple(R.Tensor, R.Tensor)) -> R.Tensor:
+ x = t[0]
+ y = t[1]
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ return w
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(t: R.Tuple(R.Tensor, R.Tensor)) -> R.Tensor:
+ with R.dataflow():
+ x = t[0]
+ y = t[1]
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ R.output(w)
+ return w
+
+
+class TestInnerFunction(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ @R.function(pure=False)
+ def inner_func(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(x, z)
+ v = R.add(y, w)
+ _ = R.print(format="oops")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.multiply(a, b)
+ return c
+
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ v = R.divide(w, z)
+ q = inner_func(w, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, a)
+ return c
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+
+ @R.function(pure=False)
+ def inner_func(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(x, z)
+ v = R.add(y, w)
+ R.output(v)
+ _ = R.print(format="oops")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.multiply(a, b)
+ R.output(c)
+ return c
+
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ v = R.divide(w, z)
+ R.output(inner_func, v, w)
+ q = inner_func(w, v)
+ with R.dataflow():
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, a)
+ R.output(c)
+ return c
+
Review Comment:
Great idea. I'll have to handle that case though.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416239066
##########
tests/python/relax/test_transform_extract_dataflow_blocks.py:
##########
@@ -0,0 +1,268 @@
+# 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 import relax
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+
+class ExtractCompare(tvm.testing.CompareBeforeAfter):
+ transform = relax.transform.ExtractDataflowBlocks()
+
+
+# functions that will not change
+class TestTrivial(ExtractCompare):
+ @I.ir_module
+ class Before:
+ # already a DF block
+ @R.function
+ def main(A: R.Tensor, B: R.Tensor):
+ with R.dataflow():
+ x = R.add(A, B)
+ y = R.multiply(x, A)
+ z = R.add(x, y)
+ q = R.multiply(y, z)
+ p = R.add(z, q)
+ R.output(p)
+ return p
+
+ # too small
+ @R.function
+ def func(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ x = R.add(A, B)
+ y = R.subtract(x, B)
+ return y
+
+ # too few pure ops between non-dataflow ops
+ @R.function(pure=False)
+ def func2(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ _ = R.print(format="Hi there!")
+ y = R.add(A, B)
+ _ = R.print(y, format="Sum: {}")
+ x = R.multiply(y, y)
+ if R.const(False):
+ _ = R.print(format="True branch")
+ q = R.add(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ else:
+ _ = R.print(format="False branch")
+ q = R.subtract(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ p = R.multiply(w, w)
+ return p
+
+ Expected = Before
+
+
+class TestBasic(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ return v
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ return v
+
+
+class TestMultipleBlocks(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ _ = R.print(format="Hi mom!")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ return d
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ _ = R.print(format="Hi mom!")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ R.output(d)
+ return d
+
+
+class TestExtractInsideBranches(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ if R.const(True):
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ else:
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, b)
+ d = R.add(c, c)
+ e = R.multiply(d, d)
+ return e
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+
+ if R.const(True):
+ with R.dataflow():
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ R.output(b)
+ # weird but the parser requires this construct
Review Comment:
Yeah, I was just hoping having `b = R.multiply(a, a)` be the last assignment would work. It's not the worst thing in the world that the parser expects a specific assignment like that, just a little weird (per the comment)
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416231743
##########
src/relax/transform/extract_dataflow_blocks.cc:
##########
@@ -0,0 +1,140 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/relax/transform/extract_dataflow_blocks.cc
+ * \brief Pass for extracting groups of pure operations without
+ * dataflow into dataflow blocks.
+ */
+
+#include <tvm/relax/expr.h>
+#include <tvm/relax/expr_functor.h>
+#include <tvm/relax/transform.h>
+#include <tvm/relax/utils.h>
+
+namespace tvm {
+namespace relax {
+
+class DataflowBlockExtractor : public ExprMutator {
+ public:
+ explicit DataflowBlockExtractor(size_t min_size) : ExprMutator(), min_size_(min_size) {}
+
+ Expr VisitExpr_(const SeqExprNode* seq) override {
+ Array<BindingBlock> new_blocks;
+ Expr new_body = VisitExpr(seq->body);
+ bool changed = !new_body.same_as(seq->body);
+ for (auto block : seq->blocks) {
+ BindingBlock new_block = this->VisitBindingBlock(block);
+ changed = changed || !new_block.same_as(block);
+ if (new_block.as<DataflowBlock>()) {
+ new_blocks.push_back(new_block);
+ continue;
+ }
+
+ // for a binding block, attempt to extract dataflow blocks inside
+ auto binding_block = Downcast<BindingBlock>(new_block);
+ bool dataflow_streak = false;
+ Array<Binding> dataflow_bindings;
+ Array<Binding> non_dataflow_bindings;
+ for (size_t i = 0; i < binding_block->bindings.size(); i++) {
+ auto binding = binding_block->bindings[i];
+ Expr value = GetBoundValue(binding);
+ // dataflow values: not an if node and not an impure call
+ bool is_dataflow =
Review Comment:
Sounds good. I was thinking that this might be an easy way to implement the folding into adjacent DataflowBlocks, since the merging is already handled there. The downside is that it wouldn't allow customization of `min_size` as part of the transform.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416183516
##########
src/relax/transform/extract_dataflow_blocks.cc:
##########
@@ -0,0 +1,140 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/relax/transform/extract_dataflow_blocks.cc
+ * \brief Pass for extracting groups of pure operations without
+ * dataflow into dataflow blocks.
+ */
+
+#include <tvm/relax/expr.h>
+#include <tvm/relax/expr_functor.h>
+#include <tvm/relax/transform.h>
+#include <tvm/relax/utils.h>
+
+namespace tvm {
+namespace relax {
+
+class DataflowBlockExtractor : public ExprMutator {
+ public:
+ explicit DataflowBlockExtractor(size_t min_size) : ExprMutator(), min_size_(min_size) {}
+
+ Expr VisitExpr_(const SeqExprNode* seq) override {
+ Array<BindingBlock> new_blocks;
+ Expr new_body = VisitExpr(seq->body);
+ bool changed = !new_body.same_as(seq->body);
+ for (auto block : seq->blocks) {
+ BindingBlock new_block = this->VisitBindingBlock(block);
+ changed = changed || !new_block.same_as(block);
+ if (new_block.as<DataflowBlock>()) {
+ new_blocks.push_back(new_block);
+ continue;
+ }
+
+ // for a binding block, attempt to extract dataflow blocks inside
+ auto binding_block = Downcast<BindingBlock>(new_block);
+ bool dataflow_streak = false;
+ Array<Binding> dataflow_bindings;
+ Array<Binding> non_dataflow_bindings;
+ for (size_t i = 0; i < binding_block->bindings.size(); i++) {
+ auto binding = binding_block->bindings[i];
+ Expr value = GetBoundValue(binding);
+ // dataflow values: not an if node and not an impure call
+ bool is_dataflow =
+ (!value.as<IfNode>()) && (!(value.as<CallNode>() && IsImpureCall(Downcast<Call>(value))));
+ if (!dataflow_streak) {
+ // we can start a dataflow streak
+ if (is_dataflow) {
+ dataflow_streak = true;
+ dataflow_bindings = {binding};
+ } else {
+ non_dataflow_bindings.push_back(binding);
+ }
+ } else {
+ if (is_dataflow) {
+ // extend the streak
+ dataflow_bindings.push_back(binding);
+ } else {
+ // this is the end of the streak
+ dataflow_streak = false;
+
+ // if the df block is below the minimum length, combine the blocks
+ // and reset the dataflow collection
+ if (dataflow_bindings.size() < min_size_) {
+ non_dataflow_bindings.insert(non_dataflow_bindings.end(), dataflow_bindings.begin(),
Review Comment:
That's a good point. Currently it's merged into the following (non-dataflow) block. This is a case worth tracking, I'll think about how to get it.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416180814
##########
src/relax/transform/canonicalize_bindings.cc:
##########
@@ -118,9 +119,28 @@ class CanonicalizePlanner : public ExprVisitor {
}
private:
+ void VisitExpr_(const FunctionNode* func) override {
+ // for functions, treat any free vars as used outside their home DF block
Review Comment:
In the specification, I actually wrote that we should not allow DF vars to be captured at all. Unless there's a compelling need to deal with that extra complexity, I'm happy to stick with that.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416230395
##########
src/relax/transform/canonicalize_bindings.cc:
##########
@@ -118,9 +119,28 @@ class CanonicalizePlanner : public ExprVisitor {
}
private:
+ void VisitExpr_(const FunctionNode* func) override {
+ // for functions, treat any free vars as used outside their home DF block
Review Comment:
That sounds like a plan, and I agree that it's not currently worth the extra complexity.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky merged PR #16204:
URL: https://github.com/apache/tvm/pull/16204
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1415875372
##########
tests/python/relax/test_transform_extract_dataflow_blocks.py:
##########
@@ -0,0 +1,268 @@
+# 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 import relax
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+
+class ExtractCompare(tvm.testing.CompareBeforeAfter):
+ transform = relax.transform.ExtractDataflowBlocks()
+
+
+# functions that will not change
+class TestTrivial(ExtractCompare):
+ @I.ir_module
+ class Before:
+ # already a DF block
+ @R.function
+ def main(A: R.Tensor, B: R.Tensor):
+ with R.dataflow():
+ x = R.add(A, B)
+ y = R.multiply(x, A)
+ z = R.add(x, y)
+ q = R.multiply(y, z)
+ p = R.add(z, q)
+ R.output(p)
+ return p
+
+ # too small
+ @R.function
+ def func(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ x = R.add(A, B)
+ y = R.subtract(x, B)
+ return y
+
+ # too few pure ops between non-dataflow ops
+ @R.function(pure=False)
+ def func2(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ _ = R.print(format="Hi there!")
+ y = R.add(A, B)
+ _ = R.print(y, format="Sum: {}")
+ x = R.multiply(y, y)
+ if R.const(False):
+ _ = R.print(format="True branch")
+ q = R.add(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ else:
+ _ = R.print(format="False branch")
+ q = R.subtract(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ p = R.multiply(w, w)
+ return p
+
+ Expected = Before
+
+
+class TestBasic(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ return v
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ return v
+
+
+class TestMultipleBlocks(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ _ = R.print(format="Hi mom!")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ return d
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ _ = R.print(format="Hi mom!")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ R.output(d)
+ return d
+
+
+class TestExtractInsideBranches(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ if R.const(True):
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ else:
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, b)
+ d = R.add(c, c)
+ e = R.multiply(d, d)
+ return e
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+
+ if R.const(True):
+ with R.dataflow():
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ R.output(b)
+ # weird but the parser requires this construct
Review Comment:
Definitely agreed that it is weird. I think it's because the relax if/else is an expression, not a statement, so the last assignment ends up being used as the binding of the `relax::IfNode` outside of the dataflow block.
##########
tests/python/relax/test_transform_extract_dataflow_blocks.py:
##########
@@ -0,0 +1,268 @@
+# 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 import relax
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+
+class ExtractCompare(tvm.testing.CompareBeforeAfter):
+ transform = relax.transform.ExtractDataflowBlocks()
+
+
+# functions that will not change
+class TestTrivial(ExtractCompare):
+ @I.ir_module
+ class Before:
+ # already a DF block
+ @R.function
+ def main(A: R.Tensor, B: R.Tensor):
+ with R.dataflow():
+ x = R.add(A, B)
+ y = R.multiply(x, A)
+ z = R.add(x, y)
+ q = R.multiply(y, z)
+ p = R.add(z, q)
+ R.output(p)
+ return p
+
+ # too small
+ @R.function
+ def func(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ x = R.add(A, B)
+ y = R.subtract(x, B)
+ return y
+
+ # too few pure ops between non-dataflow ops
+ @R.function(pure=False)
+ def func2(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ _ = R.print(format="Hi there!")
+ y = R.add(A, B)
+ _ = R.print(y, format="Sum: {}")
+ x = R.multiply(y, y)
+ if R.const(False):
+ _ = R.print(format="True branch")
+ q = R.add(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ else:
+ _ = R.print(format="False branch")
+ q = R.subtract(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ p = R.multiply(w, w)
+ return p
+
+ Expected = Before
+
+
+class TestBasic(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ return v
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ return v
+
+
+class TestMultipleBlocks(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ _ = R.print(format="Hi mom!")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ return d
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ _ = R.print(format="Hi mom!")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ R.output(d)
+ return d
+
+
+class TestExtractInsideBranches(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ if R.const(True):
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ else:
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, b)
+ d = R.add(c, c)
+ e = R.multiply(d, d)
+ return e
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+
+ if R.const(True):
+ with R.dataflow():
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ R.output(b)
+ # weird but the parser requires this construct
+ c = b
+ else:
+ with R.dataflow():
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ R.output(b)
+ c = b
+ with R.dataflow():
+ d = R.multiply(c, c)
+ e = R.add(d, d)
+ f = R.multiply(e, e)
+ R.output(f)
+ return f
+
+
+class TestTreatNonCallAsPure(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(t: R.Tuple(R.Tensor, R.Tensor)) -> R.Tensor:
+ x = t[0]
Review Comment:
Can we also add a test cases for other non-call expressions in addition to `TupleGetItem`? I think `Tuple`, `ShapeTuple`, `Constant`, and `PrimValue` would be the main ones.
##########
src/relax/transform/extract_dataflow_blocks.cc:
##########
@@ -0,0 +1,140 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/relax/transform/extract_dataflow_blocks.cc
+ * \brief Pass for extracting groups of pure operations without
+ * dataflow into dataflow blocks.
+ */
+
+#include <tvm/relax/expr.h>
+#include <tvm/relax/expr_functor.h>
+#include <tvm/relax/transform.h>
+#include <tvm/relax/utils.h>
+
+namespace tvm {
+namespace relax {
+
+class DataflowBlockExtractor : public ExprMutator {
+ public:
+ explicit DataflowBlockExtractor(size_t min_size) : ExprMutator(), min_size_(min_size) {}
+
+ Expr VisitExpr_(const SeqExprNode* seq) override {
+ Array<BindingBlock> new_blocks;
+ Expr new_body = VisitExpr(seq->body);
+ bool changed = !new_body.same_as(seq->body);
+ for (auto block : seq->blocks) {
+ BindingBlock new_block = this->VisitBindingBlock(block);
+ changed = changed || !new_block.same_as(block);
+ if (new_block.as<DataflowBlock>()) {
Review Comment:
Would be good to comment that we don't need to explicitly merge a new dataflow block with a preceding dataflow block, as that is handled during `IRNormalizer`.
##########
src/relax/transform/extract_dataflow_blocks.cc:
##########
@@ -0,0 +1,140 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/relax/transform/extract_dataflow_blocks.cc
+ * \brief Pass for extracting groups of pure operations without
+ * dataflow into dataflow blocks.
+ */
+
+#include <tvm/relax/expr.h>
+#include <tvm/relax/expr_functor.h>
+#include <tvm/relax/transform.h>
+#include <tvm/relax/utils.h>
+
+namespace tvm {
+namespace relax {
+
+class DataflowBlockExtractor : public ExprMutator {
+ public:
+ explicit DataflowBlockExtractor(size_t min_size) : ExprMutator(), min_size_(min_size) {}
+
+ Expr VisitExpr_(const SeqExprNode* seq) override {
+ Array<BindingBlock> new_blocks;
+ Expr new_body = VisitExpr(seq->body);
+ bool changed = !new_body.same_as(seq->body);
+ for (auto block : seq->blocks) {
+ BindingBlock new_block = this->VisitBindingBlock(block);
+ changed = changed || !new_block.same_as(block);
+ if (new_block.as<DataflowBlock>()) {
+ new_blocks.push_back(new_block);
+ continue;
+ }
+
+ // for a binding block, attempt to extract dataflow blocks inside
+ auto binding_block = Downcast<BindingBlock>(new_block);
+ bool dataflow_streak = false;
+ Array<Binding> dataflow_bindings;
+ Array<Binding> non_dataflow_bindings;
+ for (size_t i = 0; i < binding_block->bindings.size(); i++) {
+ auto binding = binding_block->bindings[i];
+ Expr value = GetBoundValue(binding);
+ // dataflow values: not an if node and not an impure call
+ bool is_dataflow =
+ (!value.as<IfNode>()) && (!(value.as<CallNode>() && IsImpureCall(Downcast<Call>(value))));
+ if (!dataflow_streak) {
+ // we can start a dataflow streak
+ if (is_dataflow) {
+ dataflow_streak = true;
+ dataflow_bindings = {binding};
+ } else {
+ non_dataflow_bindings.push_back(binding);
+ }
+ } else {
+ if (is_dataflow) {
+ // extend the streak
+ dataflow_bindings.push_back(binding);
+ } else {
+ // this is the end of the streak
+ dataflow_streak = false;
+
+ // if the df block is below the minimum length, combine the blocks
+ // and reset the dataflow collection
+ if (dataflow_bindings.size() < min_size_) {
+ non_dataflow_bindings.insert(non_dataflow_bindings.end(), dataflow_bindings.begin(),
Review Comment:
What happens if `dataflow_bindings.size() < min_size_`, but the preceding block was also a dataflow block? As a user, I'd want it to be merged into the preceding block anyways, but it looks like it would be left as-is.
##########
src/relax/transform/canonicalize_bindings.cc:
##########
@@ -118,9 +119,28 @@ class CanonicalizePlanner : public ExprVisitor {
}
private:
+ void VisitExpr_(const FunctionNode* func) override {
+ // for functions, treat any free vars as used outside their home DF block
+ bool cache = inside_dataflow_;
+ inside_dataflow_ = false;
+ auto free_vars = FreeVars(GetRef<Function>(func));
+ for (auto var : free_vars) {
+ used_outside_home_dataflow_.insert(var);
+ }
+ ExprVisitor::VisitExpr_(func);
+ inside_dataflow_ = cache;
+ }
+
+
+ void VisitBindingBlock_(const BindingBlockNode* block) override {
+ current_block_ = GetRef<BindingBlock>(block);
+ ExprVisitor::VisitBindingBlock_(block);
+ }
+
void VisitBindingBlock_(const DataflowBlockNode* block) override {
bool cache = inside_dataflow_;
inside_dataflow_ = true;
+ current_block_ = GetRef<DataflowBlock>(block);
Review Comment:
Can we validate that `current_block_` is `NullOpt` before setting it, and explicitly reset `current_block_` to `NullOpt` after visiting the block? As it is, `current_block_` is still set when visiting the body of a `SeqExpr`, which could lead to accidental incorrect usage in the future.
##########
src/relax/transform/canonicalize_bindings.cc:
##########
@@ -118,9 +119,28 @@ class CanonicalizePlanner : public ExprVisitor {
}
private:
+ void VisitExpr_(const FunctionNode* func) override {
+ // for functions, treat any free vars as used outside their home DF block
Review Comment:
I think this definition makes sense for inner functions, but an edge case came to mind. If the inner function is defined within a `DataflowBlock`, and the only enclosed variables it uses are from defined within the same dataflow block as the inner function, should this usage be counted as being "outside the home dataflow"?
I don't know if this is an edge case worth handling, but it did come to mind.
##########
src/relax/transform/canonicalize_bindings.cc:
##########
@@ -170,17 +191,26 @@ class CanonicalizePlanner : public ExprVisitor {
}
void VisitExpr_(const VarNode* var) override {
- if (!inside_dataflow_) {
- used_outside_dataflow_.insert(GetRef<Var>(var));
+ auto var_ref = GetRef<Var>(var);
+ // if a var is used in a dataflow block but *not* the one
+ // where it was defined, it also needs to be exposed, so also we treat that as
+ // used outside of a dataflow block
+ if (!inside_dataflow_ ||
+ (def_blocks_.count(var_ref) && !current_block_.same_as(def_blocks_.at(var_ref)))) {
+ used_outside_home_dataflow_.insert(GetRef<Var>(var));
}
}
bool inside_dataflow_{false};
+ BindingBlock current_block_;
Review Comment:
Since `current_block_` is default constructed and may be empty, can we change the type to `Optional<BindingBlock>`?
##########
src/relax/transform/extract_dataflow_blocks.cc:
##########
@@ -0,0 +1,140 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/relax/transform/extract_dataflow_blocks.cc
+ * \brief Pass for extracting groups of pure operations without
+ * dataflow into dataflow blocks.
+ */
+
+#include <tvm/relax/expr.h>
+#include <tvm/relax/expr_functor.h>
+#include <tvm/relax/transform.h>
+#include <tvm/relax/utils.h>
+
+namespace tvm {
+namespace relax {
+
+class DataflowBlockExtractor : public ExprMutator {
+ public:
+ explicit DataflowBlockExtractor(size_t min_size) : ExprMutator(), min_size_(min_size) {}
+
+ Expr VisitExpr_(const SeqExprNode* seq) override {
+ Array<BindingBlock> new_blocks;
+ Expr new_body = VisitExpr(seq->body);
+ bool changed = !new_body.same_as(seq->body);
+ for (auto block : seq->blocks) {
+ BindingBlock new_block = this->VisitBindingBlock(block);
+ changed = changed || !new_block.same_as(block);
+ if (new_block.as<DataflowBlock>()) {
+ new_blocks.push_back(new_block);
+ continue;
+ }
+
+ // for a binding block, attempt to extract dataflow blocks inside
+ auto binding_block = Downcast<BindingBlock>(new_block);
+ bool dataflow_streak = false;
+ Array<Binding> dataflow_bindings;
+ Array<Binding> non_dataflow_bindings;
+ for (size_t i = 0; i < binding_block->bindings.size(); i++) {
+ auto binding = binding_block->bindings[i];
+ Expr value = GetBoundValue(binding);
+ // dataflow values: not an if node and not an impure call
+ bool is_dataflow =
Review Comment:
This might be a bit cheeky of a way to implement the functionality, but what if we just made a `DataflowBlock` for every bound value that satisfies the `is_dataflow` check? The `IRNormalizer` combines adjacent blocks together, so the adjacent dataflow blocks would all get merged together at that point.
##########
python/tvm/relax/transform/transform.py:
##########
@@ -262,6 +262,23 @@ def LambdaLift() -> tvm.ir.transform.Pass:
return _ffi_api.LambdaLift()
+def ExtractDataflowBlocks(min_size: int = 3) -> tvm.ir.transform.Pass:
Review Comment:
Should we have the default min size be `2` instead? Having a minimum size makes sense, as a dataflow block with size of 1 doesn't have any impact, but a size of 2 does provide semantic information for downstream passes looking for a dataflow block.
##########
tests/python/relax/test_transform_extract_dataflow_blocks.py:
##########
@@ -0,0 +1,268 @@
+# 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 import relax
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+
+class ExtractCompare(tvm.testing.CompareBeforeAfter):
+ transform = relax.transform.ExtractDataflowBlocks()
+
+
+# functions that will not change
+class TestTrivial(ExtractCompare):
+ @I.ir_module
+ class Before:
+ # already a DF block
+ @R.function
+ def main(A: R.Tensor, B: R.Tensor):
+ with R.dataflow():
+ x = R.add(A, B)
+ y = R.multiply(x, A)
+ z = R.add(x, y)
+ q = R.multiply(y, z)
+ p = R.add(z, q)
+ R.output(p)
+ return p
+
+ # too small
+ @R.function
+ def func(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ x = R.add(A, B)
+ y = R.subtract(x, B)
+ return y
+
+ # too few pure ops between non-dataflow ops
+ @R.function(pure=False)
+ def func2(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ _ = R.print(format="Hi there!")
+ y = R.add(A, B)
+ _ = R.print(y, format="Sum: {}")
+ x = R.multiply(y, y)
+ if R.const(False):
+ _ = R.print(format="True branch")
+ q = R.add(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ else:
+ _ = R.print(format="False branch")
+ q = R.subtract(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ p = R.multiply(w, w)
+ return p
+
+ Expected = Before
+
+
+class TestBasic(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ return v
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ return v
+
+
+class TestMultipleBlocks(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ _ = R.print(format="Hi mom!")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ return d
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ _ = R.print(format="Hi mom!")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ R.output(d)
+ return d
+
+
+class TestExtractInsideBranches(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ if R.const(True):
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ else:
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, b)
+ d = R.add(c, c)
+ e = R.multiply(d, d)
+ return e
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+
+ if R.const(True):
+ with R.dataflow():
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ R.output(b)
+ # weird but the parser requires this construct
+ c = b
+ else:
+ with R.dataflow():
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ R.output(b)
+ c = b
+ with R.dataflow():
+ d = R.multiply(c, c)
+ e = R.add(d, d)
+ f = R.multiply(e, e)
+ R.output(f)
+ return f
+
+
+class TestTreatNonCallAsPure(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(t: R.Tuple(R.Tensor, R.Tensor)) -> R.Tensor:
+ x = t[0]
+ y = t[1]
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ return w
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(t: R.Tuple(R.Tensor, R.Tensor)) -> R.Tensor:
+ with R.dataflow():
+ x = t[0]
+ y = t[1]
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ R.output(w)
+ return w
+
+
+class TestInnerFunction(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
Review Comment:
This test validates a call to an inner impure function, but I'm not seeing any tests that validate the behavior when calling an inner pure function, or a pure/impure function contained within the same IRModule. Can we add three additional variants of this test to cover these cases?
##########
tests/python/relax/test_transform_extract_dataflow_blocks.py:
##########
@@ -0,0 +1,268 @@
+# 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 import relax
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+
+class ExtractCompare(tvm.testing.CompareBeforeAfter):
+ transform = relax.transform.ExtractDataflowBlocks()
+
+
+# functions that will not change
+class TestTrivial(ExtractCompare):
+ @I.ir_module
+ class Before:
+ # already a DF block
+ @R.function
+ def main(A: R.Tensor, B: R.Tensor):
+ with R.dataflow():
+ x = R.add(A, B)
+ y = R.multiply(x, A)
+ z = R.add(x, y)
+ q = R.multiply(y, z)
+ p = R.add(z, q)
+ R.output(p)
+ return p
+
+ # too small
+ @R.function
+ def func(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ x = R.add(A, B)
+ y = R.subtract(x, B)
+ return y
+
+ # too few pure ops between non-dataflow ops
+ @R.function(pure=False)
+ def func2(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ _ = R.print(format="Hi there!")
+ y = R.add(A, B)
+ _ = R.print(y, format="Sum: {}")
+ x = R.multiply(y, y)
+ if R.const(False):
+ _ = R.print(format="True branch")
+ q = R.add(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ else:
+ _ = R.print(format="False branch")
+ q = R.subtract(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ p = R.multiply(w, w)
+ return p
+
+ Expected = Before
+
+
+class TestBasic(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ return v
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ return v
+
+
+class TestMultipleBlocks(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ _ = R.print(format="Hi mom!")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ return d
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ _ = R.print(format="Hi mom!")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ R.output(d)
+ return d
+
+
+class TestExtractInsideBranches(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ if R.const(True):
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ else:
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, b)
+ d = R.add(c, c)
+ e = R.multiply(d, d)
+ return e
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+
+ if R.const(True):
+ with R.dataflow():
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ R.output(b)
+ # weird but the parser requires this construct
+ c = b
+ else:
+ with R.dataflow():
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ R.output(b)
+ c = b
+ with R.dataflow():
+ d = R.multiply(c, c)
+ e = R.add(d, d)
+ f = R.multiply(e, e)
+ R.output(f)
+ return f
+
+
+class TestTreatNonCallAsPure(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(t: R.Tuple(R.Tensor, R.Tensor)) -> R.Tensor:
+ x = t[0]
+ y = t[1]
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ return w
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(t: R.Tuple(R.Tensor, R.Tensor)) -> R.Tensor:
+ with R.dataflow():
+ x = t[0]
+ y = t[1]
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ R.output(w)
+ return w
+
+
+class TestInnerFunction(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ @R.function(pure=False)
+ def inner_func(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(x, z)
+ v = R.add(y, w)
+ _ = R.print(format="oops")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.multiply(a, b)
+ return c
+
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ v = R.divide(w, z)
+ q = inner_func(w, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, a)
+ return c
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+
+ @R.function(pure=False)
+ def inner_func(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(x, z)
+ v = R.add(y, w)
+ R.output(v)
+ _ = R.print(format="oops")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.multiply(a, b)
+ R.output(c)
+ return c
+
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ v = R.divide(w, z)
+ R.output(inner_func, v, w)
+ q = inner_func(w, v)
+ with R.dataflow():
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, a)
+ R.output(c)
+ return c
+
Review Comment:
Can we add two test cases for making dataflow blocks when the individual streak doesn't satisfy the `min_size_` requirement, but the presence of a dataflow block next to the potentially-dataflow binding still avoids making a tiny dataflow block.
```python
class TestMergeWithPrecedingDataflowBlock(ExtractCompare):
@I.ir_module
class Before:
@R.function
def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
with R.dataflow():
z = R.add(x, y)
w = R.multiply(z, y)
R.output(w)
# The single binding of `v = R.add` would normally not be
# enough to make a dataflow block, as `1 < min_size == 3`.
v = R.add(w, x)
return v
@I.ir_module
class Expected:
@R.function
def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
with R.dataflow():
z = R.add(x, y)
w = R.multiply(z, y)
# However, it occurs just after an existing dataflow
# block, and can be merged into it.
v = R.add(w, x)
R.output(v)
return v
class TestMergeWithNextDataflowBlock(ExtractCompare):
@I.ir_module
class Before:
@R.function
def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
# The single binding of `z = R.add` would normally not be
# enough to make a dataflow block, as `1 < min_size == 3`.
z = R.add(x, y)
# However, it occurs just before an existing dataflow
# block, and can be merged into it.
with R.dataflow():
w = R.multiply(z, y)
v = R.add(w, x)
R.output(v)
return v
@I.ir_module
class Expected:
@R.function
def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
with R.dataflow():
z = R.add(x, y)
w = R.multiply(z, y)
v = R.add(w, x)
R.output(v)
return v
```
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416274971
##########
tests/python/relax/test_transform_extract_dataflow_blocks.py:
##########
@@ -0,0 +1,268 @@
+# 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 import relax
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+
+class ExtractCompare(tvm.testing.CompareBeforeAfter):
+ transform = relax.transform.ExtractDataflowBlocks()
+
+
+# functions that will not change
+class TestTrivial(ExtractCompare):
+ @I.ir_module
+ class Before:
+ # already a DF block
+ @R.function
+ def main(A: R.Tensor, B: R.Tensor):
+ with R.dataflow():
+ x = R.add(A, B)
+ y = R.multiply(x, A)
+ z = R.add(x, y)
+ q = R.multiply(y, z)
+ p = R.add(z, q)
+ R.output(p)
+ return p
+
+ # too small
+ @R.function
+ def func(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ x = R.add(A, B)
+ y = R.subtract(x, B)
+ return y
+
+ # too few pure ops between non-dataflow ops
+ @R.function(pure=False)
+ def func2(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ _ = R.print(format="Hi there!")
+ y = R.add(A, B)
+ _ = R.print(y, format="Sum: {}")
+ x = R.multiply(y, y)
+ if R.const(False):
+ _ = R.print(format="True branch")
+ q = R.add(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ else:
+ _ = R.print(format="False branch")
+ q = R.subtract(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ p = R.multiply(w, w)
+ return p
+
+ Expected = Before
+
+
+class TestBasic(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ return v
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ return v
+
+
+class TestMultipleBlocks(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ _ = R.print(format="Hi mom!")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ return d
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ _ = R.print(format="Hi mom!")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ R.output(d)
+ return d
+
+
+class TestExtractInsideBranches(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ if R.const(True):
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ else:
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, b)
+ d = R.add(c, c)
+ e = R.multiply(d, d)
+ return e
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+
+ if R.const(True):
+ with R.dataflow():
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ R.output(b)
+ # weird but the parser requires this construct
+ c = b
+ else:
+ with R.dataflow():
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ R.output(b)
+ c = b
+ with R.dataflow():
+ d = R.multiply(c, c)
+ e = R.add(d, d)
+ f = R.multiply(e, e)
+ R.output(f)
+ return f
+
+
+class TestTreatNonCallAsPure(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(t: R.Tuple(R.Tensor, R.Tensor)) -> R.Tensor:
+ x = t[0]
+ y = t[1]
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ return w
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(t: R.Tuple(R.Tensor, R.Tensor)) -> R.Tensor:
+ with R.dataflow():
+ x = t[0]
+ y = t[1]
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ R.output(w)
+ return w
+
+
+class TestInnerFunction(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ @R.function(pure=False)
+ def inner_func(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(x, z)
+ v = R.add(y, w)
+ _ = R.print(format="oops")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.multiply(a, b)
+ return c
+
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ v = R.divide(w, z)
+ q = inner_func(w, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, a)
+ return c
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+
+ @R.function(pure=False)
+ def inner_func(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(x, z)
+ v = R.add(y, w)
+ R.output(v)
+ _ = R.print(format="oops")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.multiply(a, b)
+ R.output(c)
+ return c
+
+ z = R.add(x, y)
+ w = R.multiply(z, z)
+ v = R.divide(w, z)
+ R.output(inner_func, v, w)
+ q = inner_func(w, v)
+ with R.dataflow():
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, a)
+ R.output(c)
+ return c
+
Review Comment:
Done, thank you
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "slyubomirsky (via GitHub)" <gi...@apache.org>.
slyubomirsky commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416185052
##########
src/relax/transform/extract_dataflow_blocks.cc:
##########
@@ -0,0 +1,140 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/relax/transform/extract_dataflow_blocks.cc
+ * \brief Pass for extracting groups of pure operations without
+ * dataflow into dataflow blocks.
+ */
+
+#include <tvm/relax/expr.h>
+#include <tvm/relax/expr_functor.h>
+#include <tvm/relax/transform.h>
+#include <tvm/relax/utils.h>
+
+namespace tvm {
+namespace relax {
+
+class DataflowBlockExtractor : public ExprMutator {
+ public:
+ explicit DataflowBlockExtractor(size_t min_size) : ExprMutator(), min_size_(min_size) {}
+
+ Expr VisitExpr_(const SeqExprNode* seq) override {
+ Array<BindingBlock> new_blocks;
+ Expr new_body = VisitExpr(seq->body);
+ bool changed = !new_body.same_as(seq->body);
+ for (auto block : seq->blocks) {
+ BindingBlock new_block = this->VisitBindingBlock(block);
+ changed = changed || !new_block.same_as(block);
+ if (new_block.as<DataflowBlock>()) {
+ new_blocks.push_back(new_block);
+ continue;
+ }
+
+ // for a binding block, attempt to extract dataflow blocks inside
+ auto binding_block = Downcast<BindingBlock>(new_block);
+ bool dataflow_streak = false;
+ Array<Binding> dataflow_bindings;
+ Array<Binding> non_dataflow_bindings;
+ for (size_t i = 0; i < binding_block->bindings.size(); i++) {
+ auto binding = binding_block->bindings[i];
+ Expr value = GetBoundValue(binding);
+ // dataflow values: not an if node and not an impure call
+ bool is_dataflow =
Review Comment:
Cheeky indeed. I'll keep it as-is for now but if we run into trouble, that's a good option haha.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [Unity][Transform] Pass for automatically extracting DataflowBlocks [tvm]
Posted by "Lunderberg (via GitHub)" <gi...@apache.org>.
Lunderberg commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416625283
##########
tests/python/relax/test_transform_convert_dataflow.py:
##########
@@ -0,0 +1,493 @@
+# 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 import relax
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+
+class ExtractCompare(tvm.testing.CompareBeforeAfter):
+ transform = relax.transform.ConvertToDataflow()
+
+
+# functions that will not change
+class TestTrivial(ExtractCompare):
+ @I.ir_module
+ class Before:
+ # already a DF block
+ @R.function
+ def main(A: R.Tensor, B: R.Tensor):
+ with R.dataflow():
+ x = R.add(A, B)
+ y = R.multiply(x, A)
+ z = R.add(x, y)
+ q = R.multiply(y, z)
+ p = R.add(z, q)
+ R.output(p)
+ return p
+
+ # too small
+ @R.function
+ def func(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ x = R.add(A, B)
+ return x
+
+ # too few pure ops between non-dataflow ops
+ @R.function(pure=False)
+ def func2(A: R.Tensor, B: R.Tensor) -> R.Tensor:
+ _ = R.print(format="Hi there!")
+ y = R.add(A, B)
+ _ = R.print(y, format="Sum: {}")
+ x = R.multiply(y, y)
+ if R.const(False):
+ _ = R.print(format="True branch")
+ q = R.add(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ else:
+ _ = R.print(format="False branch")
+ q = R.subtract(x, y)
+ _ = R.print(q, format="Value of q: {}")
+ w = q
+ p = R.multiply(w, w)
+ return p
+
+ Expected = Before
+
+
+class TestBasic(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ return v
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ return v
+
+
+class TestMultipleBlocks(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ _ = R.print(format="Hi mom!")
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ return d
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+ _ = R.print(format="Hi mom!")
+ with R.dataflow():
+ a = R.multiply(v, v)
+ b = R.add(a, a)
+ c = R.subtract(b, a)
+ d = R.add(c, c)
+ R.output(d)
+ return d
+
+
+class TestExtractInsideBranches(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ if R.const(True):
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ else:
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ c = R.multiply(b, b)
+ d = R.add(c, c)
+ e = R.multiply(d, d)
+ return e
+
+ @I.ir_module
+ class Expected:
+ @R.function(pure=False)
+ def main(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ with R.dataflow():
+ z = R.add(x, y)
+ w = R.multiply(z, y)
+ v = R.add(w, x)
+ R.output(v)
+
+ if R.const(True):
+ with R.dataflow():
+ q = R.multiply(v, v)
+ a = R.add(q, q)
+ b = R.multiply(a, a)
+ R.output(b)
+ # weird but the parser requires this construct
+ c = b
+ else:
+ with R.dataflow():
+ q = R.add(v, v)
+ a = R.multiply(q, q)
+ b = R.add(a, a)
+ R.output(b)
+ c = b
+ with R.dataflow():
+ d = R.multiply(c, c)
+ e = R.add(d, d)
+ f = R.multiply(e, e)
+ R.output(f)
+ return f
+
+
+class TestTreatNonCallAsPure(ExtractCompare):
+ @I.ir_module
+ class Before:
+ @R.function
+ def tuples_and_const(x: R.Tensor, y: R.Tensor) -> R.Tensor:
+ t1 = (x, y, x)
+ t2 = (y, y, x)
+ c = R.const([1, 2, 3], dtype="int32")
+ return c
+
+ @R.function
+ def shapes() -> R.Shape:
+ s1 = R.shape((1, 2, 3))
+ s2 = R.shape((4, 5, 6))
+ s3 = R.shape((7, 8, 9))
+ return s3
+
+ @R.function
+ def prim_values():
Review Comment:
I think a more stable way to have the pretty-printing from TVMScript would be to remove all of the try/catch that only exist to re-throw as DiagnosticError, and instead allow the exception to bubble up to the entry point. Catching the exception at the entry point, the `exc.__traceback__` can then be inspected. The TVMScript location of the error is then the lowest stack frame that contains a variable of type `doc.AST`, and that object can be provided to `Parser.report_error` as the location of the error in the TVMScript.
This would avoid requiring every node to have its own try/catch/re-throw, and avoid the segfault that we ran into here.
@junrushao Since you've spent more time on the TVMScript parser than I have, what are your thoughts on this idea?
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org