Function.cpp

#include <atomic>
#include <cstdlib>
#include <memory>
#include <set>
#include <utility>

#include "CSE.h"
#include "Func.h"
#include "Function.h"
#include "IR.h"
#include "IREquality.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "IRPrinter.h"
#include "ParallelRVar.h"
#include "Random.h"
#include "Scope.h"
#include "Var.h"

namespace Halide {
namespace Internal {

using std::map;
using std::pair;
using std::string;
using std::vector;

typedef map<FunctionPtr, FunctionPtr> DeepCopyMap;

struct FunctionContents;

namespace {

// Weaken all the references to a particular Function to break
// reference cycles. Also count the number of references found.
class WeakenFunctionPtrs : public IRMutator {
    using IRMutator::visit;

    Expr visit(const Call *c) override {
        Expr expr = IRMutator::visit(c);
        c = expr.as<Call>();
        internal_assert(c);
        if (c->func.defined() &&
            c->func.get() == func) {
            FunctionPtr ptr = c->func;
            ptr.weaken();
            expr = Call::make(c->type, c->name, c->args, c->call_type,
                              ptr, c->value_index,
                              c->image, c->param);
            count++;
        }
        return expr;
    }
    FunctionContents *func;

public:
    int count = 0;
    WeakenFunctionPtrs(FunctionContents *f)
        : func(f) {
    }
};

}  // namespace

struct FunctionContents {
    std::string name;
    std::string origin_name;
    std::vector<Type> output_types;

    /** Optional type constraints on the Function:
     * - If empty, there are no constraints.
     * - If size == 1, the Func is only allowed to have values of Expr with that type
     * - If size > 1, the Func is only allowed to have values of Tuple with those types
     *
     * Note that when this is nonempty, then output_types should match
     * required_types for all defined Functions.
     */
    std::vector<Type> required_types;

    /** Optional dimension constraints on the Function:
     * - If required_dims == AnyDims, there are no constraints.
     * - Otherwise, the Function's dimensionality must exactly match required_dims.
     */
    int required_dims = AnyDims;

    // The names of the dimensions of the Function. Corresponds to the
    // LHS of the pure definition if there is one. Is also the initial
    // stage of the dims and storage_dims. Used to identify dimensions
    // of the Function by name.
    std::vector<string> args;

    // Function-specific schedule. This schedule is applied to all stages
    // within the function.
    FuncSchedule func_schedule;

    Definition init_def;
    std::vector<Definition> updates;

    std::string debug_file;

    std::vector<Parameter> output_buffers;

    std::vector<ExternFuncArgument> extern_arguments;
    std::string extern_function_name;

    NameMangling extern_mangling = NameMangling::Default;
    DeviceAPI extern_function_device_api = DeviceAPI::Host;
    Expr extern_proxy_expr;

    bool trace_loads = false, trace_stores = false, trace_realizations = false;
    std::vector<string> trace_tags;

    bool no_profiling = false;

    bool frozen = false;

    void accept(IRVisitor *visitor) const {
        func_schedule.accept(visitor);

        if (init_def.defined()) {
            init_def.accept(visitor);
        }
        for (const Definition &def : updates) {
            def.accept(visitor);
        }

        if (!extern_function_name.empty()) {
            for (const ExternFuncArgument &i : extern_arguments) {
                if (i.is_func()) {
                    user_assert(i.func.get() != this)
                        << "Extern Func has itself as an argument";
                    i.func->accept(visitor);
                } else if (i.is_expr()) {
                    i.expr.accept(visitor);
                }
            }
            if (extern_proxy_expr.defined()) {
                extern_proxy_expr.accept(visitor);
            }
        }

        for (const Parameter &i : output_buffers) {
            for (size_t j = 0; j < args.size(); j++) {
                if (i.min_constraint(j).defined()) {
                    i.min_constraint(j).accept(visitor);
                }
                if (i.stride_constraint(j).defined()) {
                    i.stride_constraint(j).accept(visitor);
                }
                if (i.extent_constraint(j).defined()) {
                    i.extent_constraint(j).accept(visitor);
                }
            }
        }
    }

    // Pass an IRMutator through to all Exprs referenced in the FunctionContents
    void mutate(IRMutator *mutator) {
        func_schedule.mutate(mutator);

        if (init_def.defined()) {
            init_def.mutate(mutator);
        }
        for (Definition &def : updates) {
            def.mutate(mutator);
        }

        if (!extern_function_name.empty()) {
            for (ExternFuncArgument &i : extern_arguments) {
                if (i.is_expr()) {
                    i.expr = mutator->mutate(i.expr);
                }
            }
            extern_proxy_expr = mutator->mutate(extern_proxy_expr);
        }
    }
};

struct FunctionGroup {
    mutable RefCount ref_count;
    vector<FunctionContents> members;
};

FunctionContents *FunctionPtr::get() const {
    return &(group()->members[idx]);
}

template<>
RefCount &ref_count<FunctionGroup>(const FunctionGroup *f) noexcept {
    return f->ref_count;
}

template<>
void destroy<FunctionGroup>(const FunctionGroup *f) {
    delete f;
}

namespace {

// All variables present in any part of a function definition must
// either be pure args, elements of the reduction domain, parameters
// (i.e. attached to some Parameter object), or part of a let node
// internal to the expression
struct CheckVars : public IRGraphVisitor {
    vector<string> pure_args;
    ReductionDomain reduction_domain;
    Scope<> defined_internally;
    const std::string name;
    bool unbound_reduction_vars_ok = false;
    bool pure;

    CheckVars(const std::string &n, bool pure)
        : name(n), pure(pure) {
    }

    using IRVisitor::visit;

    void visit(const Let *let) override {
        let->value.accept(this);
        ScopedBinding<> bind(defined_internally, let->name);
        let->body.accept(this);
    }

    void visit(const Call *op) override {
        IRGraphVisitor::visit(op);
        bool proper_func = (name == op->name || op->call_type != Call::Halide);
        if (op->call_type == Call::Halide &&
            op->func.defined() &&
            op->name != name) {
            Function func = Function(op->func);
            proper_func = (name == op->name || func.has_pure_definition() || func.has_extern_definition());
        }
        if (pure) {
            user_assert(proper_func)
                << "In pure definition of Func \"" << name << "\":\n"
                << "Can't call Func \"" << op->name
                << "\" because it has not yet been defined,"
                << " and it is not a recursive call.\n";
        } else {
            user_assert(proper_func)
                << "In update definition of Func \"" << name << "\":\n"
                << "Can't call Func \"" << op->name
                << "\" because it has not yet been defined.\n";
            if (op->name == name && op->call_type == Call::Halide) {
                for (size_t i = 0; i < op->args.size(); i++) {
                    const Variable *var = op->args[i].as<Variable>();
                    if (!pure_args[i].empty()) {
                        user_assert(var && var->name == pure_args[i])
                            << "In update definition of Func \"" << name << "\":\n"
                            << "All of a function's recursive references to itself"
                            << " in update definitions must contain the same pure"
                            << " variables in the same places as on the left-hand-side.\n";
                    }
                }
            }
        }
    }

    void visit(const Variable *var) override {
        // Is it a parameter?
        if (var->param.defined()) {
            return;
        }

        // Was it defined internally by a let expression?
        if (defined_internally.contains(var->name)) {
            return;
        }

        // Is it a pure argument?
        for (auto &pure_arg : pure_args) {
            if (var->name == pure_arg) {
                return;
            }
        }

        // Is it in a reduction domain?
        if (var->reduction_domain.defined()) {
            if (!reduction_domain.defined()) {
                reduction_domain = var->reduction_domain;
                return;
            } else if (var->reduction_domain.same_as(reduction_domain)) {
                // It's in a reduction domain we already know about
                return;
            } else {
                user_error << "Multiple reduction domains found in definition of Func \"" << name << "\"\n";
            }
        } else if (reduction_domain.defined() && unbound_reduction_vars_ok) {
            // Is it one of the RVars from the reduction domain we already
            // know about (this can happen in the RDom predicate).
            for (const ReductionVariable &rv : reduction_domain.domain()) {
                if (rv.var == var->name) {
                    return;
                }
            }
        }

        user_error << "Undefined variable \"" << var->name << "\" in definition of Func \"" << name << "\"\n";
    }
};

// Mark all functions found in an expr as frozen.
class FreezeFunctions : public IRGraphVisitor {
    using IRGraphVisitor::visit;

    const string &func;

    void visit(const Call *op) override {
        IRGraphVisitor::visit(op);
        if (op->call_type == Call::Halide &&
            op->func.defined() &&
            op->name != func) {
            Function f(op->func);
            f.freeze();
        }
    }

public:
    FreezeFunctions(const string &f)
        : func(f) {
    }
};

}  // namespace

// A counter to use in tagging random variables.
// Note that this will be reset by Internal::reset_random_counters().
std::atomic<int> random_variable_counter = 0;

Function::Function(const FunctionPtr &ptr)
    : contents(ptr) {
    contents.strengthen();
    internal_assert(ptr.defined())
        << "Can't construct Function from undefined FunctionContents ptr\n";
}

Function::Function(const std::string &n) {
    for (size_t i = 0; i < n.size(); i++) {
        user_assert(n[i] != '.')
            << "Func name \"" << n << "\" is invalid. "
            << "Func names may not contain the character '.', "
            << "as it is used internally by Halide as a separator\n";
    }
    contents.strong = new FunctionGroup;
    contents.strong->members.resize(1);
    contents->name = n;
    contents->origin_name = n;
}

Function::Function(const std::vector<Type> &required_types, int required_dims, const std::string &n)
    : Function(n) {
    user_assert(required_dims >= AnyDims);
    contents->required_types = required_types;
    contents->required_dims = required_dims;
}

void Function::update_with_deserialization(const std::string &name,
                                           const std::string &origin_name,
                                           const std::vector<Halide::Type> &output_types,
                                           const std::vector<Halide::Type> &required_types,
                                           int required_dims,
                                           const std::vector<std::string> &args,
                                           const FuncSchedule &func_schedule,
                                           const Definition &init_def,
                                           const std::vector<Definition> &updates,
                                           const std::string &debug_file,
                                           const std::vector<Parameter> &output_buffers,
                                           const std::vector<ExternFuncArgument> &extern_arguments,
                                           const std::string &extern_function_name,
                                           NameMangling name_mangling,
                                           DeviceAPI device_api,
                                           const Expr &extern_proxy_expr,
                                           bool trace_loads,
                                           bool trace_stores,
                                           bool trace_realizations,
                                           const std::vector<std::string> &trace_tags,
                                           bool no_profiling,
                                           bool frozen) {
    contents->name = name;
    contents->origin_name = origin_name;
    contents->output_types = output_types;
    contents->required_types = required_types;
    contents->required_dims = required_dims;
    contents->args = args;
    contents->func_schedule = func_schedule;
    contents->init_def = init_def;
    contents->updates = updates;
    contents->debug_file = debug_file;
    contents->output_buffers = output_buffers;
    contents->extern_arguments = extern_arguments;
    contents->extern_function_name = extern_function_name;
    contents->extern_mangling = name_mangling;
    contents->extern_function_device_api = device_api;
    contents->extern_proxy_expr = extern_proxy_expr;
    contents->trace_loads = trace_loads;
    contents->trace_stores = trace_stores;
    contents->trace_realizations = trace_realizations;
    contents->trace_tags = trace_tags;
    contents->no_profiling = no_profiling;
    contents->frozen = frozen;
}

namespace {

template<typename T>
struct PrintTypeList {
    const std::vector<T> &list_;

    explicit PrintTypeList(const std::vector<T> &list)
        : list_(list) {
    }

    friend std::ostream &operator<<(std::ostream &s, const PrintTypeList &self) {
        const size_t n = self.list_.size();
        if (n != 1) {
            s << "(";
        }
        const char *comma = "";
        for (const auto &t : self.list_) {
            if constexpr (std::is_same_v<Type, T>) {
                s << comma << t;
            } else {
                s << comma << t.type();
            }
            comma = ", ";
        }
        if (n != 1) {
            s << ")";
        }
        return s;
    }
};

bool types_match(const std::vector<Type> &types, const std::vector<Expr> &exprs) {
    size_t n = types.size();
    if (n != exprs.size()) {
        return false;
    }
    for (size_t i = 0; i < n; i++) {
        if (types[i] != exprs[i].type()) {
            return false;
        }
    }
    return true;
}

}  // namespace

void Function::check_types(const Expr &e) const {
    check_types(std::vector<Expr>{e});
}

void Function::check_types(const Tuple &t) const {
    check_types(t.as_vector());
}

void Function::check_types(const Type &t) const {
    check_types(std::vector<Type>{t});
}

void Function::check_types(const std::vector<Expr> &exprs) const {
    if (!contents->required_types.empty()) {
        user_assert(types_match(contents->required_types, exprs))
            << "Func \"" << name() << "\" is constrained to only hold values of type " << PrintTypeList(contents->required_types)
            << " but is defined with values of type " << PrintTypeList(exprs) << ".\n";
    }
}

void Function::check_types(const std::vector<Type> &types) const {
    if (!contents->required_types.empty()) {
        user_assert(contents->required_types == types)
            << "Func \"" << name() << "\" is constrained to only hold values of type " << PrintTypeList(contents->required_types)
            << " but is defined with values of type " << PrintTypeList(types) << ".\n";
    }
}

void Function::check_dims(int dims) const {
    if (contents->required_dims != AnyDims) {
        user_assert(contents->required_dims == dims)
            << "Func \"" << name() << "\" is constrained to have exactly " << contents->required_dims
            << " dimensions, but is defined with " << dims << " dimensions.\n";
    }
}

namespace {

// Return deep-copy of ExternFuncArgument 'src'
ExternFuncArgument deep_copy_extern_func_argument_helper(const ExternFuncArgument &src,
                                                         DeepCopyMap &copied_map) {
    ExternFuncArgument copy;
    copy.arg_type = src.arg_type;
    copy.buffer = src.buffer;
    copy.expr = src.expr;
    copy.image_param = src.image_param;

    if (!src.func.defined()) {  // No need to deep-copy the func if it's undefined
        internal_assert(!src.is_func())
            << "ExternFuncArgument has type FuncArg but has no function definition\n";
        return copy;
    }

    // If the FunctionContents has already been deep-copied previously, i.e.
    // it's in the 'copied_map', use the deep-copied version from the map instead
    // of creating a new deep-copy
    FunctionPtr &copied_func = copied_map[src.func];
    internal_assert(copied_func.defined());
    copy.func = copied_func;
    return copy;
}

}  // namespace

void Function::deep_copy(const FunctionPtr &copy, DeepCopyMap &copied_map) const {
    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.
    copied_map[contents] = copy;

    debug(4) << "Deep-copy function contents: \"" << contents->name << "\"\n";

    copy->name = contents->name;
    copy->origin_name = contents->origin_name;
    copy->args = contents->args;
    copy->output_types = contents->output_types;
    copy->debug_file = contents->debug_file;
    copy->extern_function_name = contents->extern_function_name;
    copy->extern_mangling = contents->extern_mangling;
    copy->extern_function_device_api = contents->extern_function_device_api;
    copy->extern_proxy_expr = contents->extern_proxy_expr;
    copy->trace_loads = contents->trace_loads;
    copy->trace_stores = contents->trace_stores;
    copy->trace_realizations = contents->trace_realizations;
    copy->trace_tags = contents->trace_tags;
    copy->no_profiling = contents->no_profiling;
    copy->frozen = contents->frozen;
    copy->output_buffers = contents->output_buffers;
    copy->func_schedule = contents->func_schedule.deep_copy(copied_map);

    // Copy the pure definition
    if (contents->init_def.defined()) {
        copy->init_def = contents->init_def.get_copy();
        internal_assert(copy->init_def.is_init());
        internal_assert(copy->init_def.schedule().rvars().empty())
            << "Init definition shouldn't have reduction domain\n";
    }

    for (const Definition &def : contents->updates) {
        internal_assert(!def.is_init());
        Definition def_copy = def.get_copy();
        internal_assert(!def_copy.is_init());
        copy->updates.push_back(std::move(def_copy));
    }

    for (const ExternFuncArgument &e : contents->extern_arguments) {
        ExternFuncArgument e_copy = deep_copy_extern_func_argument_helper(e, copied_map);
        copy->extern_arguments.push_back(std::move(e_copy));
    }
}

void Function::deep_copy(string name, const FunctionPtr &copy, DeepCopyMap &copied_map) const {
    deep_copy(copy, copied_map);
    copy->name = std::move(name);
}

void Function::define(const vector<string> &args, vector<Expr> values) {
    user_assert(!frozen())
        << "Func " << name() << " cannot be given a new pure definition, "
        << "because it has already been realized or used in the definition of another Func.\n";
    user_assert(!has_extern_definition())
        << "In pure definition of Func \"" << name() << "\":\n"
        << "Func with extern definition cannot be given a pure definition.\n";
    user_assert(!name().empty()) << "A Func may not have an empty name.\n";
    for (auto &value : values) {
        user_assert(value.defined())
            << "In pure definition of Func \"" << name() << "\":\n"
            << "Undefined expression in right-hand-side of definition.\n";
    }

    // Make sure all the vars in the value are either args or are
    // attached to some parameter
    CheckVars check(name(), true);
    check.pure_args = args;
    for (const auto &value : values) {
        value.accept(&check);
    }

    // Freeze all called functions
    FreezeFunctions freezer(name());
    // TODO: Check for calls to undefined Funcs
    for (const auto &value : values) {
        value.accept(&freezer);
    }

    // Make sure all the vars in the args have unique non-empty names
    for (size_t i = 0; i < args.size(); i++) {
        user_assert(!args[i].empty())
            << "In pure definition of Func \"" << name() << "\":\n"
            << "In left-hand-side of definition, argument "
            << i << " has an empty name.\n";
        for (size_t j = 0; j < i; j++) {
            user_assert(args[i] != args[j])
                << "In pure definition of Func \"" << name() << "\":\n"
                << "In left-hand-side of definition, arguments "
                << i << " and " << j
                << " both have the name \"" + args[i] + "\"\n";
        }
    }

    for (auto &value : values) {
        value = common_subexpression_elimination(value);
    }

    // Tag calls to random() with the free vars
    int tag = random_variable_counter++;
    vector<VarOrRVar> free_vars;
    free_vars.reserve(args.size());
    for (const auto &arg : args) {
        free_vars.emplace_back(Var(arg));
    }
    for (auto &value : values) {
        value = lower_random(value, free_vars, tag);
    }

    user_assert(!check.reduction_domain.defined())
        << "In pure definition of Func \"" << name() << "\":\n"
        << "Reduction domain referenced in pure function definition.\n";

    if (!contents.defined()) {
        contents.strong = new FunctionGroup;
        contents.strong->members.resize(1);
        contents->name = unique_name('f');
        contents->origin_name = contents->name;
    }

    user_assert(!contents->init_def.defined())
        << "In pure definition of Func \"" << name() << "\":\n"
        << "Func is already defined.\n";

    check_types(values);
    check_dims((int)args.size());
    contents->args = args;

    std::vector<Expr> init_def_args;
    init_def_args.resize(args.size());
    for (size_t i = 0; i < args.size(); i++) {
        init_def_args[i] = Var(args[i]);
    }

    // If the function is inductive,
    // the value and args might refer back to the
    // function itself, introducing circular references and hence
    // memory leaks. We need to break these cycles.
    WeakenFunctionPtrs weakener(contents.get());
    for (auto &arg : init_def_args) {
        arg = weakener.mutate(arg);
    }
    for (auto &value : values) {
        value = weakener.mutate(value);
    }
    if (check.reduction_domain.defined()) {
        check.reduction_domain.set_predicate(
            weakener.mutate(check.reduction_domain.predicate()));
    }

    ReductionDomain rdom;
    contents->init_def = Definition(init_def_args, values, rdom, true);

    for (const auto &arg : args) {
        DimType dtype = DimType::PureVar;
        if (is_inductive(arg)) {
            dtype = DimType::InductiveVar;
        }
        Dim d = {arg, ForType::Serial, DeviceAPI::None, dtype};
        contents->init_def.schedule().dims().push_back(d);
        StorageDim sd = {arg};
        contents->func_schedule.storage_dims().push_back(sd);
    }

    // Add the dummy outermost dim
    {
        Dim d = {Var::outermost().name(), ForType::Serial, DeviceAPI::None, DimType::PureVar};
        contents->init_def.schedule().dims().push_back(d);
    }

    contents->output_types.resize(values.size());
    for (size_t i = 0; i < contents->output_types.size(); i++) {
        contents->output_types[i] = values[i].type();
    }

    if (!contents->required_types.empty()) {
        // Just a reality check; mismatches here really should have been caught earlier
        internal_assert(contents->required_types == contents->output_types);
    }
    if (contents->required_dims != AnyDims) {
        // Just a reality check; mismatches here really should have been caught earlier
        internal_assert(contents->required_dims == (int)args.size());
    }

    if (contents->output_buffers.empty()) {
        create_output_buffers(contents->output_types, (int)args.size());
    }
}

void Function::create_output_buffers(const std::vector<Type> &types, int dims) const {
    internal_assert(contents->output_buffers.empty());
    internal_assert(!types.empty() && dims != AnyDims);

    for (size_t i = 0; i < types.size(); i++) {
        string buffer_name = name();
        if (types.size() > 1) {
            buffer_name += '.' + std::to_string((int)i);
        }
        Parameter output(types[i], true, dims, buffer_name);
        contents->output_buffers.push_back(output);
    }
}

void Function::define_update(const vector<Expr> &_args, vector<Expr> values, const ReductionDomain &rdom) {
    int update_idx = static_cast<int>(contents->updates.size());

    user_assert(!name().empty())
        << "Func has an empty name.\n";
    user_assert(has_pure_definition())
        << "In update definition " << update_idx << " of Func \"" << name() << "\":\n"
        << "Can't add an update definition without a pure definition first.\n";
    user_assert(!frozen())
        << "Func " << name() << " cannot be given a new update definition, "
        << "because it has already been realized or used in the definition of another Func.\n";
    user_assert(!is_inductive())
        << "In update definition " << update_idx << " of Func \"" << name() << "\":\n"
        << "Inductive functions cannot have update definitions.\n";

    for (auto &value : values) {
        user_assert(value.defined())
            << "In update definition " << update_idx << " of Func \"" << name() << "\":\n"
            << "Undefined expression in right-hand-side of update.\n";
    }

    // Check the dimensionality matches
    user_assert((int)_args.size() == dimensions())
        << "In update definition " << update_idx << " of Func \"" << name() << "\":\n"
        << "Dimensionality of update definition must match dimensionality of pure definition.\n";

    user_assert(values.size() == contents->init_def.values().size())
        << "In update definition " << update_idx << " of Func \"" << name() << "\":\n"
        << "Number of tuple elements for update definition must "
        << "match number of tuple elements for pure definition.\n";

    const auto &pure_def_vals = contents->init_def.values();
    for (size_t i = 0; i < values.size(); i++) {
        // Check that pure value and the update value have the same
        // type.  Without this check, allocations may be the wrong size
        // relative to what update code expects.
        Type pure_type = pure_def_vals[i].type();
        if (pure_type != values[i].type()) {
            std::ostringstream err;
            err << "In update definition " << update_idx << " of Func \"" << name() << "\":\n";
            if (!values.empty()) {
                err << "Tuple element " << i << " of update definition has type ";
            } else {
                err << "Update definition has type ";
            }
            err << values[i].type() << ", but pure definition has type " << pure_type;
            user_error << err.str() << "\n";
        }
        values[i] = common_subexpression_elimination(values[i]);
    }

    vector<Expr> args(_args.size());
    for (size_t i = 0; i < args.size(); i++) {
        args[i] = common_subexpression_elimination(_args[i]);
    }

    // The pure args are those naked vars in the args that are not in
    // a reduction domain and are not parameters and line up with the
    // pure args in the pure definition.
    bool pure = true;
    vector<string> pure_args(args.size());
    for (size_t i = 0; i < args.size(); i++) {
        pure_args[i] = "";  // Will never match a var name
        user_assert(args[i].defined())
            << "In update definition " << update_idx << " of Func \"" << name() << "\":\n"
            << "Argument " << i
            << " in left-hand-side of update definition is undefined.\n";
        if (const Variable *var = args[i].as<Variable>()) {
            if (!var->param.defined() &&
                !var->reduction_domain.defined() &&
                var->name == contents->args[i]) {
                pure_args[i] = var->name;
            } else {
                pure = false;
            }
        } else {
            pure = false;
        }
    }

    // Make sure all the vars in the args and the value are either
    // pure args, in the reduction domain, or a parameter. Also checks
    // that recursive references to the function contain all the pure
    // vars in the LHS in the correct places.
    CheckVars check(name(), false);
    check.pure_args = pure_args;
    for (const auto &arg : args) {
        arg.accept(&check);
    }
    for (const auto &value : values) {
        value.accept(&check);
    }
    if (!check.reduction_domain.defined()) {
        // Use the provided one
        check.reduction_domain = rdom;
    } else if (rdom.defined()) {
        // This is an internal error because the ability to pass an explicit
        // RDom is not exposed to the front-end. At the time of writing this is
        // only used by rfactor.
        internal_assert(rdom.same_as(check.reduction_domain))
            << "In update definition " << update_idx << " of Func \"" << name() << "\":\n"
            << "Explicit reduction domain passed to Function::define_update, "
            << "but another reduction domain was referred to by the args or values.\n"
            << "Explicit reduction domain passed:\n"
            << RDom(rdom) << "\n"
            << "Found reduction domain:\n"
            << RDom(check.reduction_domain) << "\n";
    }

    if (check.reduction_domain.defined()) {
        check.unbound_reduction_vars_ok = true;
        check.reduction_domain.predicate().accept(&check);
    }

    // Freeze all called functions
    FreezeFunctions freezer(name());
    for (const auto &arg : args) {
        arg.accept(&freezer);
    }
    for (const auto &value : values) {
        value.accept(&freezer);
    }

    // Freeze the reduction domain if defined
    if (check.reduction_domain.defined()) {
        check.reduction_domain.predicate().accept(&freezer);
        check.reduction_domain.freeze();
    }

    // Tag calls to random() with the free vars
    vector<VarOrRVar> free_vars;
    int num_free_vars = (int)pure_args.size();
    if (check.reduction_domain.defined()) {
        num_free_vars += (int)check.reduction_domain.domain().size();
    }
    free_vars.reserve(num_free_vars);
    for (const auto &pure_arg : pure_args) {
        if (!pure_arg.empty()) {
            free_vars.emplace_back(Var(pure_arg));
        }
    }
    if (check.reduction_domain.defined()) {
        for (size_t i = 0; i < check.reduction_domain.domain().size(); i++) {
            free_vars.emplace_back(RVar(check.reduction_domain, i));
        }
    }
    int tag = random_variable_counter++;
    for (auto &arg : args) {
        arg = lower_random(arg, free_vars, tag);
    }
    for (auto &value : values) {
        value = lower_random(value, free_vars, tag);
    }
    if (check.reduction_domain.defined()) {
        check.reduction_domain.set_predicate(lower_random(check.reduction_domain.predicate(), free_vars, tag));
    }

    // The update value and args probably refer back to the
    // function itself, introducing circular references and hence
    // memory leaks. We need to break these cycles.
    WeakenFunctionPtrs weakener(contents.get());
    for (auto &arg : args) {
        arg = weakener.mutate(arg);
    }
    for (auto &value : values) {
        value = weakener.mutate(value);
    }
    if (check.reduction_domain.defined()) {
        check.reduction_domain.set_predicate(
            weakener.mutate(check.reduction_domain.predicate()));
    }

    Definition r(args, values, check.reduction_domain, false);
    internal_assert(!r.is_init()) << "Should have been an update definition\n";

    // First add any reduction domain
    if (check.reduction_domain.defined()) {
        for (const auto &rvar : check.reduction_domain.domain()) {
            // Is this RVar actually pure (safe to parallelize and
            // reorder)? It's pure if one value of the RVar can never
            // access from the same memory that another RVar is
            // writing to.
            const string &v = rvar.var;

            bool pure = can_parallelize_rvar(v, name(), r);
            Dim d = {v, ForType::Serial, DeviceAPI::None,
                     pure ? DimType::PureRVar : DimType::ImpureRVar};
            r.schedule().dims().push_back(d);
        }
    }

    // Then add the pure args outside of that
    for (const auto &pure_arg : pure_args) {
        if (!pure_arg.empty()) {
            Dim d = {pure_arg, ForType::Serial, DeviceAPI::None, DimType::PureVar};
            r.schedule().dims().push_back(d);
        }
    }

    // Then the dummy outermost dim
    {
        Dim d = {Var::outermost().name(), ForType::Serial, DeviceAPI::None, DimType::PureVar};
        r.schedule().dims().push_back(d);
    }

    // If there's no recursive reference, no reduction domain, and all
    // the args are pure, then this definition completely hides
    // earlier ones!
    if (!check.reduction_domain.defined() &&
        weakener.count == 0 &&
        pure) {
        user_warning
            << "In update definition " << update_idx << " of Func \"" << name() << "\":\n"
            << "Update definition completely hides earlier definitions, "
            << " because all the arguments are pure, it contains no self-references, "
            << " and no reduction domain. This may be an accidental re-definition of "
            << " an already-defined function.\n";
    }

    contents->updates.push_back(r);
}

void Function::define_extern(const std::string &function_name,
                             const std::vector<ExternFuncArgument> &extern_args,
                             const std::vector<Type> &types,
                             const std::vector<Var> &args,
                             NameMangling mangling,
                             DeviceAPI device_api) {
    check_types(types);
    check_dims((int)args.size());

    user_assert(!has_pure_definition() && !has_update_definition())
        << "In extern definition for Func \"" << name() << "\":\n"
        << "Func with a pure definition cannot have an extern definition.\n";

    user_assert(!has_extern_definition())
        << "In extern definition for Func \"" << name() << "\":\n"
        << "Func already has an extern definition.\n";

    std::vector<string> arg_names;
    std::vector<Expr> arg_exprs;
    for (const auto &arg : args) {
        arg_names.push_back(arg.name());
        arg_exprs.push_back(arg);
    }
    contents->args = arg_names;
    contents->extern_function_name = function_name;
    contents->extern_arguments = extern_args;
    contents->output_types = types;
    contents->extern_mangling = mangling;
    contents->extern_function_device_api = device_api;

    std::vector<Expr> values;
    contents->output_buffers.clear();
    for (size_t i = 0; i < types.size(); i++) {
        string buffer_name = name();
        if (types.size() > 1) {
            buffer_name += '.' + std::to_string((int)i);
        }
        Parameter output(types[i], true, (int)args.size(), buffer_name);
        contents->output_buffers.push_back(output);

        values.push_back(undef(types[i]));
    }

    contents->init_def = Definition(arg_exprs, values, ReductionDomain(), true);

    // Reset the storage dims to match the pure args
    contents->func_schedule.storage_dims().clear();
    contents->init_def.schedule().dims().clear();
    for (size_t i = 0; i < args.size(); i++) {
        StorageDim sd = {arg_names[i]};
        contents->func_schedule.storage_dims().push_back(sd);
        Dim d = {arg_names[i], ForType::Extern, DeviceAPI::None, DimType::PureVar};
        contents->init_def.schedule().dims().push_back(d);
    }
    // Add the dummy outermost dim
    Dim d = {Var::outermost().name(), ForType::Serial, DeviceAPI::None, DimType::PureVar};
    contents->init_def.schedule().dims().push_back(d);
}

void Function::accept(IRVisitor *visitor) const {
    contents->accept(visitor);
}

void Function::mutate(IRMutator *mutator) {
    contents->mutate(mutator);
}

const std::string &Function::name() const {
    return contents->name;
}

const std::string &Function::origin_name() const {
    return contents->origin_name;
}

Definition &Function::definition() {
    internal_assert(contents->init_def.defined());
    return contents->init_def;
}

const Definition &Function::definition() const {
    internal_assert(contents->init_def.defined());
    return contents->init_def;
}

const std::vector<std::string> &Function::args() const {
    return contents->args;
}

bool Function::is_pure_arg(const std::string &name) const {
    return std::find(args().begin(), args().end(), name) != args().end();
}

int Function::dimensions() const {
    return (int)args().size();
}

int Function::outputs() const {
    return (int)output_types().size();
}

const std::vector<Type> &Function::output_types() const {
    return contents->output_types;
}

const std::vector<Type> &Function::required_types() const {
    return contents->required_types;
}

int Function::required_dimensions() const {
    return contents->required_dims;
}

const std::vector<Expr> &Function::values() const {
    static const std::vector<Expr> empty;
    if (has_pure_definition()) {
        return contents->init_def.values();
    } else {
        return empty;
    }
}

FuncSchedule &Function::schedule() {
    return contents->func_schedule;
}

const FuncSchedule &Function::schedule() const {
    return contents->func_schedule;
}

const std::vector<Parameter> &Function::output_buffers() const {
    if (!contents->output_buffers.empty()) {
        return contents->output_buffers;
    }

    // If types and dims are already specified, we can go ahead and create
    // the output buffer(s) even if the Function has no pure definition yet.
    if (!contents->required_types.empty() && contents->required_dims != AnyDims) {
        create_output_buffers(contents->required_types, contents->required_dims);
        return contents->output_buffers;
    }

    user_error << "Can't access output buffer(s) of undefined Func \"" << name() << "\".\n";
    return contents->output_buffers;
}

StageSchedule &Function::update_schedule(int idx) {
    internal_assert(idx < (int)contents->updates.size()) << "Invalid update definition index\n";
    return contents->updates[idx].schedule();
}

Definition &Function::update(int idx) {
    internal_assert(idx < (int)contents->updates.size()) << "Invalid update definition index\n";
    return contents->updates[idx];
}

const Definition &Function::update(int idx) const {
    internal_assert(idx < (int)contents->updates.size()) << "Invalid update definition index\n";
    return contents->updates[idx];
}

const std::vector<Definition> &Function::updates() const {
    return contents->updates;
}

bool Function::has_pure_definition() const {
    return contents->init_def.defined();
}

bool Function::is_inductive() const {
    class RecursiveHelper : public IRVisitor {
        using IRVisitor::visit;
        const string &func;
        void visit(const Call *op) override {
            if (op->name == func) {
                recursive = true;
            }
            IRVisitor::visit(op);
        }

    public:
        bool recursive = false;
        RecursiveHelper(const string &func)
            : func(func) {
        }
    };

    if (!has_pure_definition()) {
        return false;
    }

    RecursiveHelper r(name());
    for (const Expr &e : definition().values()) {
        e.accept(&r);
    }

    return r.recursive;
}

bool Function::is_inductive(const string &var) const {
    class RecursiveHelper : public IRVisitor {
        using IRVisitor::visit;
        const string &func;
        const string &var;
        const int &pos;
        void visit(const Call *op) override {
            if (op->name == func) {
                recursive = true;
                if (const auto &v = op->args[pos].as<Variable>()) {
                    if (v->name != var) {
                        inductive_in_var = true;
                    }
                } else {
                    inductive_in_var = true;
                }
            }
            IRVisitor::visit(op);
        }

    public:
        bool recursive = false;
        bool inductive_in_var = false;
        RecursiveHelper(const string &func, const string &var, const int &pos)
            : func(func), var(var), pos(pos) {
        }
    };

    if (!has_pure_definition()) {
        return false;
    }

    int pos = -1;
    for (size_t i = 0; i < definition().args().size(); i++) {
        if (const auto &v = definition().args()[i].as<Variable>()) {
            if (v->name == var) {
                pos = i;
            }
        }
    }
    if (pos == -1) {
        return false;
    }
    RecursiveHelper r(name(), var, pos);
    for (const Expr &e : definition().values()) {
        e.accept(&r);
    }

    return r.inductive_in_var;
}

bool Function::can_be_inlined() const {
    return is_pure() && definition().specializations().empty() && !is_inductive();
}

bool Function::has_update_definition() const {
    return !contents->updates.empty();
}

bool Function::has_extern_definition() const {
    return !contents->extern_function_name.empty();
}

NameMangling Function::extern_definition_name_mangling() const {
    return contents->extern_mangling;
}

Expr Function::make_call_to_extern_definition(const std::vector<Expr> &args,
                                              const Target &target) const {
    internal_assert(has_extern_definition());

    Call::CallType call_type = Call::Extern;
    switch (contents->extern_mangling) {
    case NameMangling::Default:
        call_type = (target.has_feature(Target::CPlusPlusMangling) ? Call::ExternCPlusPlus : Call::Extern);
        break;
    case NameMangling::CPlusPlus:
        call_type = Call::ExternCPlusPlus;
        break;
    case NameMangling::C:
        call_type = Call::Extern;
        break;
    }
    return Call::make(Int(32), contents->extern_function_name, args, call_type, contents);
}

Expr Function::extern_definition_proxy_expr() const {
    return contents->extern_proxy_expr;
}

Expr &Function::extern_definition_proxy_expr() {
    return contents->extern_proxy_expr;
}

const std::vector<ExternFuncArgument> &Function::extern_arguments() const {
    return contents->extern_arguments;
}

std::vector<ExternFuncArgument> &Function::extern_arguments() {
    return contents->extern_arguments;
}

const std::string &Function::extern_function_name() const {
    return contents->extern_function_name;
}

DeviceAPI Function::extern_function_device_api() const {
    return contents->extern_function_device_api;
}

const std::string &Function::debug_file() const {
    return contents->debug_file;
}

std::string &Function::debug_file() {
    return contents->debug_file;
}

Function::operator ExternFuncArgument() const {
    return ExternFuncArgument(contents);
}

void Function::trace_loads() {
    contents->trace_loads = true;
}
void Function::trace_stores() {
    contents->trace_stores = true;
}
void Function::trace_realizations() {
    contents->trace_realizations = true;
}
void Function::add_trace_tag(const std::string &trace_tag) {
    contents->trace_tags.push_back(trace_tag);
}

bool Function::is_tracing_loads() const {
    return contents->trace_loads;
}
bool Function::is_tracing_stores() const {
    return contents->trace_stores;
}
bool Function::is_tracing_realizations() const {
    return contents->trace_realizations;
}
const std::vector<std::string> &Function::get_trace_tags() const {
    return contents->trace_tags;
}

void Function::lock_loop_levels() {
    auto &schedule = contents->func_schedule;
    schedule.compute_level().lock();
    schedule.store_level().lock();
    schedule.hoist_storage_level().lock();
    // If store_level is inlined, use the compute_level instead.
    // (Note that we deliberately do *not* do the same if store_level
    // is undefined.)
    if (schedule.store_level().is_inlined()) {
        schedule.store_level() = schedule.compute_level();
    }
    if (schedule.hoist_storage_level().is_inlined()) {
        schedule.hoist_storage_level() = schedule.store_level();
    }
    if (contents->init_def.defined()) {
        contents->init_def.schedule().fuse_level().level.lock();
    }
    for (Definition &def : contents->updates) {
        internal_assert(def.defined());
        def.schedule().fuse_level().level.lock();
    }
}

void Function::do_not_profile() {
    contents->no_profiling = true;
}
bool Function::should_not_profile() const {
    return contents->no_profiling;
}

void Function::freeze() {
    contents->frozen = true;
}
bool Function::frozen() const {
    return contents->frozen;
}

const map<string, FunctionPtr> &Function::wrappers() const {
    return contents->func_schedule.wrappers();
}

Function Function::new_function_in_same_group(const std::string &f) {
    int group_size = (int)(contents.group()->members.size());
    contents.group()->members.resize(group_size + 1);
    contents.group()->members[group_size].name = f;
    FunctionPtr ptr;
    ptr.strong = contents.group();
    ptr.idx = group_size;
    return Function(ptr);
}

void Function::add_wrapper(const std::string &f, Function &wrapper) {
    wrapper.freeze();
    FunctionPtr ptr = wrapper.contents;

    // Weaken the pointer from the function to its wrapper
    ptr.weaken();
    contents->func_schedule.add_wrapper(f, ptr);

    // Weaken the pointer from the wrapper back to the function.
    WeakenFunctionPtrs weakener(contents.get());
    wrapper.mutate(&weakener);
}

const Call *Function::is_wrapper() const {
    const vector<Expr> &rhs = values();
    if (rhs.size() != 1) {
        return nullptr;
    }
    const Call *call = rhs[0].as<Call>();
    if (!call) {
        return nullptr;
    }
    vector<Expr> expected_args;
    for (const string &v : args()) {
        expected_args.push_back(Variable::make(Int(32), v));
    }
    Expr expected_rhs =
        Call::make(call->type, call->name, expected_args, call->call_type,
                   call->func, call->value_index, call->image, call->param);
    if (equal(rhs[0], expected_rhs)) {
        return call;
    } else {
        return nullptr;
    }
}

Function &Function::substitute_calls(const map<FunctionPtr, FunctionPtr> &substitutions) {
    debug(4) << "Substituting calls in " << name() << "\n";
    if (substitutions.empty()) {
        return *this;
    }

    // Replace all calls to functions listed in 'substitutions' with their wrappers.
    auto m = LambdaMutator{
        [&](auto *self, const Call *c) {
            Expr expr = self->visit_base(c);
            c = expr.as<Call>();
            internal_assert(c);

            if ((c->call_type == Call::Halide) &&
                c->func.defined() &&
                substitutions.count(c->func)) {
                auto it = substitutions.find(c->func);
                internal_assert(it != substitutions.end())
                    << "Function not in environment: " << c->func->name << "\n";
                FunctionPtr subs = it->second;
                debug(4) << "...Replace call to Func \"" << c->name << "\" with "
                         << "\"" << subs->name << "\"\n";
                expr = Call::make(c->type, subs->name, c->args, c->call_type,
                                  subs, c->value_index,
                                  c->image, c->param);
            }
            return expr;
        }};

    contents->mutate(&m);
    return *this;
}

Function &Function::substitute_calls(const Function &orig, const Function &substitute) {
    map<FunctionPtr, FunctionPtr> substitutions;
    substitutions.emplace(orig.get_contents(), substitute.get_contents());
    return substitute_calls(substitutions);
}

// Deep copy an entire Function DAG.
pair<vector<Function>, map<string, Function>> deep_copy(
    const vector<Function> &outputs, const map<string, Function> &env) {
    vector<Function> copy_outputs;
    map<string, Function> copy_env;

    // Create empty deep-copies of all Functions in 'env'
    DeepCopyMap copied_map;  // Original Function -> Deep-copy
    IntrusivePtr<FunctionGroup> group(new FunctionGroup);
    group->members.resize(env.size());
    int i = 0;
    for (const auto &iter : env) {
        // Make a weak pointer to the function to use for within-group references.
        FunctionPtr ptr;
        ptr.weak = group.get();
        ptr.idx = i;
        ptr->name = iter.second.name();
        copied_map[iter.second.get_contents()] = ptr;
        i++;
    }

    // Deep copy all Functions in 'env' into their corresponding empty copies
    for (const auto &iter : env) {
        iter.second.deep_copy(copied_map[iter.second.get_contents()], copied_map);
    }

    // Need to substitute-in all old Function references in all Exprs referenced
    // within the Function with the deep-copy versions
    for (auto &iter : copied_map) {
        Function(iter.second).substitute_calls(copied_map);
    }

    // Populate the env with the deep-copy version
    for (const auto &iter : copied_map) {
        FunctionPtr ptr = iter.second;
        copy_env.emplace(iter.first->name, Function(ptr));
    }

    for (const auto &func : outputs) {
        const auto &iter = copied_map.find(func.get_contents());
        if (iter != copied_map.end()) {
            FunctionPtr ptr = iter->second;
            debug(4) << "Adding deep-copied version to outputs: " << func.name() << "\n";
            copy_outputs.emplace_back(ptr);
        } else {
            debug(4) << "Adding original version to outputs: " << func.name() << "\n";
            copy_outputs.push_back(func);
        }
    }

    return {copy_outputs, copy_env};
}

}  // namespace Internal
}  // namespace Halide