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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
