Skip to content

Commit

Permalink
Make a deep copy of each piece of test IR so that we can parallelize
Browse files Browse the repository at this point in the history
  • Loading branch information
abadams committed Feb 23, 2024
1 parent b03b3c7 commit f8952c2
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 6 deletions.
6 changes: 4 additions & 2 deletions src/Function.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -487,8 +487,10 @@ ExternFuncArgument deep_copy_extern_func_argument_helper(const ExternFuncArgumen
} // namespace

void Function::deep_copy(const FunctionPtr &copy, DeepCopyMap &copied_map) const {
internal_assert(copy.defined() && contents.defined())
<< "Cannot deep-copy undefined Function\n";
internal_assert(copy.defined())
<< "Cannot deep-copy to undefined Function\n";
internal_assert(contents.defined())
<< "Cannot deep-copy from undefined Function\n";

// Add reference to this Function's deep-copy to the map in case of
// self-reference, e.g. self-reference in an Definition.
Expand Down
46 changes: 42 additions & 4 deletions test/correctness/simd_op_check_sve2.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -60,10 +60,6 @@ class SimdOpCheckArmSve : public SimdOpCheckTest {
return can_run_the_code;
}

bool use_multiple_threads() const override {
return false;
}

void add_tests() override {
check_arm_integer();
check_arm_float();
Expand Down Expand Up @@ -1234,6 +1230,48 @@ class SimdOpCheckArmSve : public SimdOpCheckTest {
// settings.
if (!parent.wildcard_match(parent.filter, decorated_op_name)) return;

// Create a deep copy of the expr and all Funcs referenced by it, so
// that no IR is shared between tests. This is required by the base
// class, and is why we can parallelize.
{
using namespace Halide::Internal;
class FindOutputs : public IRVisitor {
using IRVisitor::visit;
void visit(const Call *op) override {
if (op->func.defined()) {
outputs.insert(op->func);
}
IRVisitor::visit(op);
}

public:
std::set<FunctionPtr> outputs;
} finder;
e.accept(&finder);
std::vector<Function> outputs(finder.outputs.begin(), finder.outputs.end());
auto env = deep_copy(outputs, build_environment(outputs)).second;
class DeepCopy : public IRMutator {
std::map<FunctionPtr, FunctionPtr> copied;
using IRMutator::visit;
Expr visit(const Call *op) override {
if (op->func.defined()) {
auto it = env.find(op->name);
if (it != env.end()) {
return Func(it->second)(mutate(op->args));
}
}
return IRMutator::visit(op);
}
const std::map<std::string, Function> &env;

public:
DeepCopy(const std::map<std::string, Function> &env)
: env(env) {
}
} copier(env);
e = copier.mutate(e);
}

// Create Task and register
parent.tasks.emplace_back(Task{decorated_op_name, unique_name, vec_factor, e});
parent.arm_tasks.emplace(unique_name, ArmTask{std::move(instr_patterns)});
Expand Down

0 comments on commit f8952c2

Please sign in to comment.