diff --git a/src/SplitTuples.cpp b/src/SplitTuples.cpp index 686f919a25af..b2de80d12e81 100644 --- a/src/SplitTuples.cpp +++ b/src/SplitTuples.cpp @@ -69,10 +69,11 @@ class SplitTuples : public IRMutator { if (op->types.size() > 1) { // Make a nested set of realize nodes for each tuple element Stmt body = mutate(op->body); + Expr condition = mutate(op->condition); for (int i = (int)op->types.size() - 1; i >= 0; i--) { body = Realize::make(op->name + "." + std::to_string(i), {op->types[i]}, op->memory_type, - op->bounds, op->condition, body); + op->bounds, condition, body); } return body; } else { diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index 1aee6cecdd0d..62a93e25a982 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -239,6 +239,7 @@ tests(GROUPS correctness pseudostack_shares_slots.cpp python_extension_gen.cpp pytorch.cpp + realize_condition_depends_on_tuple.cpp realize_larger_than_two_gigs.cpp realize_over_shifted_domain.cpp reduction_chain.cpp diff --git a/test/correctness/realize_condition_depends_on_tuple.cpp b/test/correctness/realize_condition_depends_on_tuple.cpp new file mode 100644 index 000000000000..c4bf0b58e9a7 --- /dev/null +++ b/test/correctness/realize_condition_depends_on_tuple.cpp @@ -0,0 +1,31 @@ +#include "Halide.h" + +using namespace Halide; + +// This is a test for a bug where the condition on a realize node didn't have +// tuple-valued calls resolved if the realization was itself tuple-valued. + +int main(int argc, char **argv) { + Func f; + Param p; + f() = {p, p}; + + Func g; + g() = {4, 4}; + + Func h; + h() = g()[1]; + + // h may or may not be necessary to evaluate, depending on a load from f, + // which means g in turn may or may not be necessary to allocate. + Func out; + out() = select(f()[1] == 3, h(), 17); + + f.compute_root(); + g.compute_root(); + h.compute_root(); + out.compile_jit(); + + printf("Success!\n"); + return 0; +}