From f410b1ad92c11112007d13cdc8709486e6a7e803 Mon Sep 17 00:00:00 2001 From: Steven Johnson Date: Mon, 25 Apr 2022 17:54:38 -0700 Subject: [PATCH] Augment Halide::Func to allow for constraining Type and Dimensionality This enhances Func by allowing you to (optionally) constrain the type(s) of Exprs that the Func can contain, and/or the dimensionality of the Func. (Attempting to violate either of these will assert-fail.) There are a few goals here: - Enhanced code readability; in cases where a Func's values may not be obvious from the code flow, this can allow an in-code way of declaring it (rather than via comments) - Enhanced type enforcement; specifying constraints allows us to fail in type-mismatched compilations somewhat sooner, with somewhat better error messages. - Better symmetry for AOT/JIT code generation with ImageParam, in which the inputs (ImageParam) have a way to specify the required concrete type, but the outputs (Funcs) don't. If this is accepted, then subsequent changes will likely add uses where it makes sense (e.g., the Func associated with an ImageParam should always have both type and dimensionality specified since it will always be well-known). Note that this doesn't add any C++ template class for static declarations (e.g. `FuncT` -> `Func(Float(32), 2)`); these could be added later if desired. --- python_bindings/correctness/basics.py | 33 +++++ python_bindings/src/PyFunc.cpp | 2 + src/Buffer.h | 6 +- src/Func.cpp | 15 +++ src/Func.h | 13 ++ src/Function.cpp | 124 +++++++++++++++++- src/Function.h | 20 +++ src/Pipeline.cpp | 2 + test/error/CMakeLists.txt | 8 ++ test/error/func_expr_dim_mismatch.cpp | 17 +++ test/error/func_expr_type_mismatch.cpp | 17 +++ test/error/func_expr_update_type_mismatch.cpp | 18 +++ test/error/func_extern_dim_mismatch.cpp | 14 ++ test/error/func_extern_type_mismatch.cpp | 14 ++ test/error/func_tuple_dim_mismatch.cpp | 17 +++ test/error/func_tuple_types_mismatch.cpp | 17 +++ .../func_tuple_update_types_mismatch.cpp | 18 +++ 17 files changed, 351 insertions(+), 4 deletions(-) create mode 100644 test/error/func_expr_dim_mismatch.cpp create mode 100644 test/error/func_expr_type_mismatch.cpp create mode 100644 test/error/func_expr_update_type_mismatch.cpp create mode 100644 test/error/func_extern_dim_mismatch.cpp create mode 100644 test/error/func_extern_type_mismatch.cpp create mode 100644 test/error/func_tuple_dim_mismatch.cpp create mode 100644 test/error/func_tuple_types_mismatch.cpp create mode 100644 test/error/func_tuple_update_types_mismatch.cpp diff --git a/python_bindings/correctness/basics.py b/python_bindings/correctness/basics.py index b63bf2ff299c..86eba1c0a4f5 100644 --- a/python_bindings/correctness/basics.py +++ b/python_bindings/correctness/basics.py @@ -309,11 +309,44 @@ def test_bool_conversion(): # Verify that this doesn't fail with 'Argument passed to specialize must be of type bool' f.compute_root().specialize(True) +def test_typed_funcs(): + x = hl.Var('x') + y = hl.Var('y') + + f = hl.Func(hl.Int(32), 1, 'f') + try: + f[x, y] = hl.i32(0); + f.realize([10, 10]) + except RuntimeError as e: + assert 'is constrained to have exactly 1 dimensions, but is defined with 2 dimensions' in str(e) + else: + assert False, 'Did not see expected exception!' + + f = hl.Func(hl.Int(32), 2, 'f') + try: + f[x, y] = hl.i16(0); + f.realize([10, 10]) + except RuntimeError as e: + assert 'is constrained to only hold values of type int32 but is defined with values of type int16' in str(e) + else: + assert False, 'Did not see expected exception!' + + f = hl.Func((hl.Int(32), hl.Float(32)), 2, 'f') + try: + f[x, y] = (hl.i16(0), hl.f64(0)) + f.realize([10, 10]) + except RuntimeError as e: + assert 'is constrained to only hold values of type (int32, float32) but is defined with values of type (int16, float64)' in str(e) + else: + assert False, 'Did not see expected exception!' + + if __name__ == "__main__": test_compiletime_error() test_runtime_error() test_misused_and() test_misused_or() + test_typed_funcs() test_float_or_int() test_operator_order() test_int_promotion() diff --git a/python_bindings/src/PyFunc.cpp b/python_bindings/src/PyFunc.cpp index 8ec0f1836804..bb8e88061fd0 100644 --- a/python_bindings/src/PyFunc.cpp +++ b/python_bindings/src/PyFunc.cpp @@ -108,6 +108,8 @@ void define_func(py::module &m) { py::class_(m, "Func") .def(py::init<>()) .def(py::init()) + .def(py::init(), py::arg("required_type"), py::arg("required_dimensions"), py::arg("name")) + .def(py::init, int, std::string>(), py::arg("required_types"), py::arg("required_dimensions"), py::arg("name")) .def(py::init()) .def(py::init([](Buffer<> &b) -> Func { return Func(b); })) diff --git a/src/Buffer.h b/src/Buffer.h index eaff181f7fdc..220c009a8ea1 100644 --- a/src/Buffer.h +++ b/src/Buffer.h @@ -8,7 +8,9 @@ namespace Halide { -template +constexpr int AnyDims = Halide::Runtime::AnyDims; // -1 + +template class Buffer; struct JITUserContext; @@ -153,7 +155,7 @@ class Buffer { } public: - static constexpr int AnyDims = Halide::Runtime::AnyDims; + static constexpr int AnyDims = Halide::AnyDims; static_assert(Dims == AnyDims || Dims >= 0); typedef T ElemType; diff --git a/src/Func.cpp b/src/Func.cpp index a0c000307af6..12e25c8c0f44 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -59,6 +59,14 @@ Func::Func(const string &name) : func(unique_name(name)) { } +Func::Func(const Type &required_type, int required_dims, const string &name) + : func({required_type}, required_dims, unique_name(name)) { +} + +Func::Func(const std::vector &required_types, int required_dims, const string &name) + : func(required_types, required_dims, unique_name(name)) { +} + Func::Func() : func(make_entity_name(this, "Halide:.*:Func", 'f')) { } @@ -2926,6 +2934,8 @@ Stage FuncRef::operator=(const FuncRef &e) { } } +namespace { + // Inject a suitable base-case definition given an update // definition. This is a helper for FuncRef::operator+= and co. Func define_base_case(const Internal::Function &func, const vector &a, const Tuple &e) { @@ -2955,8 +2965,12 @@ Func define_base_case(const Internal::Function &func, const vector &a, con return define_base_case(func, a, Tuple(e)); } +} // namespace + template Stage FuncRef::func_ref_update(const Tuple &e, int init_val) { + func.check_types(e); + internal_assert(e.size() > 1); vector init_values(e.size()); @@ -2975,6 +2989,7 @@ Stage FuncRef::func_ref_update(const Tuple &e, int init_val) { template Stage FuncRef::func_ref_update(Expr e, int init_val) { + func.check_types(e); vector expanded_args = args_with_implicit_vars({e}); FuncRef self_ref = define_base_case(func, expanded_args, cast(e.type(), init_val))(expanded_args); return self_ref = BinaryOp()(Expr(self_ref), e); diff --git a/src/Func.h b/src/Func.h index 4c643ab776de..a54219100039 100644 --- a/src/Func.h +++ b/src/Func.h @@ -715,6 +715,19 @@ class Func { /** Declare a new undefined function with the given name */ explicit Func(const std::string &name); + /** Declare a new undefined function with the given name. + * The function will be constrained to represent Exprs of required_type. + * If required_dims is not AnyDims, the function will be constrained to exactly + * that many dimensions. */ + explicit Func(const Type &required_type, int required_dims, const std::string &name); + + /** Declare a new undefined function with the given name. + * If required_types is not empty, the function will be constrained to represent + * Tuples of the same arity and types. (If required_types is empty, there is no constraint.) + * If required_dims is not AnyDims, the function will be constrained to exactly + * that many dimensions. */ + explicit Func(const std::vector &required_types, int required_dims, const std::string &name); + /** Declare a new undefined function with an * automatically-generated unique name */ Func(); diff --git a/src/Function.cpp b/src/Function.cpp index f7eded59e824..7d49824cbb9a 100644 --- a/src/Function.cpp +++ b/src/Function.cpp @@ -30,6 +30,7 @@ typedef map 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 { @@ -58,6 +59,7 @@ class WeakenFunctionPtrs : public IRMutator { : func(f) { } }; + } // namespace struct FunctionContents { @@ -65,6 +67,22 @@ struct FunctionContents { std::string origin_name; std::vector 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 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 @@ -306,9 +324,100 @@ Function::Function(const std::string &n) { contents->origin_name = n; } +Function::Function(const std::vector &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; +} + +namespace { + +template +struct PrintTypeList { + const std::vector &list_; + + explicit PrintTypeList(const std::vector &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::value) { + s << comma << t; + } else { + s << comma << t.type(); + } + comma = ", "; + } + if (n != 1) { + s << ")"; + } + return s; + } +}; + +bool types_match(const std::vector &types, const std::vector &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{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{t}); +} + +void Function::check_types(const std::vector &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 &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 deep_copy_extern_func_argument_helper(const ExternFuncArgument &src, + DeepCopyMap &copied_map) { ExternFuncArgument copy; copy.arg_type = src.arg_type; copy.buffer = src.buffer; @@ -330,6 +439,8 @@ ExternFuncArgument deep_copy_extern_func_argument_helper( return copy; } +} // namespace + void Function::deep_copy(const FunctionPtr ©, DeepCopyMap &copied_map) const { internal_assert(copy.defined() && contents.defined()) << "Cannot deep-copy undefined Function\n"; @@ -456,6 +567,8 @@ void Function::define(const vector &args, vector values) { << "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 init_def_args; @@ -485,6 +598,11 @@ void Function::define(const vector &args, vector values) { 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); + } + for (size_t i = 0; i < values.size(); i++) { string buffer_name = name(); if (values.size() > 1) { @@ -703,6 +821,8 @@ void Function::define_extern(const std::string &function_name, const std::vector &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" diff --git a/src/Function.h b/src/Function.h index ce8a76ef4f17..ee67d4181a81 100644 --- a/src/Function.h +++ b/src/Function.h @@ -17,6 +17,7 @@ namespace Halide { struct ExternFuncArgument; +class Tuple; class Var; @@ -57,6 +58,13 @@ class Function { /** Construct a new function with the given name */ explicit Function(const std::string &n); + /** Construct a new function with the given name, + * with a requirement that it can only represent Expr(s) of the given type(s), + * and must have exactly the give nnumber of dimensions. + * required_types.empty() means there are no constraints on the type(s). + * required_dims == AnyDims means there are no constraints on the dimensions. */ + explicit Function(const std::vector &required_types, int required_dims, const std::string &n); + /** Construct a Function from an existing FunctionContents pointer. Must be non-null */ explicit Function(const FunctionPtr &); @@ -292,6 +300,18 @@ class Function { /** Return true iff the name matches one of the Function's pure args. */ bool is_pure_arg(const std::string &name) const; + + /** If the Function has type requirements, check that the given argument + * is compatible with them. If not, assert-fail. (If there are no type requirements, do nothing.) */ + void check_types(const Expr &e) const; + void check_types(const Tuple &t) const; + void check_types(const Type &t) const; + void check_types(const std::vector &exprs) const; + void check_types(const std::vector &types) const; + + /** If the Function has dimension requirements, check that the given argument + * is compatible with them. If not, assert-fail. (If there are no dimension requirements, do nothing.) */ + void check_dims(int dims) const; }; /** Deep copy an entire Function DAG. */ diff --git a/src/Pipeline.cpp b/src/Pipeline.cpp index 733709e2a23d..d658fb977e8e 100644 --- a/src/Pipeline.cpp +++ b/src/Pipeline.cpp @@ -687,6 +687,8 @@ Realization Pipeline::realize(JITUserContext *context, user_assert(defined()) << "Pipeline is undefined\n"; vector> bufs; for (auto &out : contents->outputs) { + user_assert((int)sizes.size() == out.dimensions()) + << "Func " << out.name() << " is defined with " << out.dimensions() << " dimensions, but realize() is requesting a realization with " << sizes.size() << " dimensions.\n"; user_assert(out.has_pure_definition() || out.has_extern_definition()) << "Can't realize Pipeline with undefined output Func: " << out.name() << ".\n"; for (Type t : out.output_types()) { bufs.emplace_back(t, nullptr, sizes); diff --git a/test/error/CMakeLists.txt b/test/error/CMakeLists.txt index ee75855332e3..063221f3bc87 100644 --- a/test/error/CMakeLists.txt +++ b/test/error/CMakeLists.txt @@ -44,6 +44,14 @@ tests(GROUPS error five_d_gpu_buffer.cpp float_arg.cpp forward_on_undefined_buffer.cpp + func_expr_dim_mismatch.cpp + func_expr_type_mismatch.cpp + func_expr_update_type_mismatch.cpp + func_extern_dim_mismatch.cpp + func_extern_type_mismatch.cpp + func_tuple_dim_mismatch.cpp + func_tuple_types_mismatch.cpp + func_tuple_update_types_mismatch.cpp implicit_args.cpp impossible_constraints.cpp init_def_should_be_all_vars.cpp diff --git a/test/error/func_expr_dim_mismatch.cpp b/test/error/func_expr_dim_mismatch.cpp new file mode 100644 index 000000000000..1218f70ecb7b --- /dev/null +++ b/test/error/func_expr_dim_mismatch.cpp @@ -0,0 +1,17 @@ +#include "Halide.h" +#include + +using namespace Halide; +using namespace Halide::Internal; + +int main(int argc, char **argv) { + Var x("x"), y("y"); + Func f(Int(32), 1, "f"); + + f(x, y) = cast(0); + + f.realize({100, 100}); + + printf("Success!\n"); + return 0; +} diff --git a/test/error/func_expr_type_mismatch.cpp b/test/error/func_expr_type_mismatch.cpp new file mode 100644 index 000000000000..1337891cacf3 --- /dev/null +++ b/test/error/func_expr_type_mismatch.cpp @@ -0,0 +1,17 @@ +#include "Halide.h" +#include + +using namespace Halide; +using namespace Halide::Internal; + +int main(int argc, char **argv) { + Var x("x"), y("y"); + Func f(Float(32), 1, "f"); + + f(x, y) = cast(0); + + f.realize({100, 100}); + + printf("Success!\n"); + return 0; +} diff --git a/test/error/func_expr_update_type_mismatch.cpp b/test/error/func_expr_update_type_mismatch.cpp new file mode 100644 index 000000000000..a26146936561 --- /dev/null +++ b/test/error/func_expr_update_type_mismatch.cpp @@ -0,0 +1,18 @@ +#include "Halide.h" +#include + +using namespace Halide; +using namespace Halide::Internal; + +int main(int argc, char **argv) { + Var x("x"), y("y"); + Func f(Float(32), 2, "f"); + + f(x, y) = 0.f; + f(x, y) += cast(0); + + f.realize({100, 100}); + + printf("Success!\n"); + return 0; +} diff --git a/test/error/func_extern_dim_mismatch.cpp b/test/error/func_extern_dim_mismatch.cpp new file mode 100644 index 000000000000..41c39ca47f9c --- /dev/null +++ b/test/error/func_extern_dim_mismatch.cpp @@ -0,0 +1,14 @@ +#include "Halide.h" +#include + +using namespace Halide; +using namespace Halide::Internal; + +int main(int argc, char **argv) { + Var x("x"), y("y"); + Func f(Float(32), 1, "f"); + f.define_extern("test", {}, Float(32), {x, y}); + f.realize({100, 100}); + printf("Success!\n"); + return 0; +} diff --git a/test/error/func_extern_type_mismatch.cpp b/test/error/func_extern_type_mismatch.cpp new file mode 100644 index 000000000000..ad137f40aca3 --- /dev/null +++ b/test/error/func_extern_type_mismatch.cpp @@ -0,0 +1,14 @@ +#include "Halide.h" +#include + +using namespace Halide; +using namespace Halide::Internal; + +int main(int argc, char **argv) { + Var x("x"), y("y"); + Func f({UInt(8), Float(64)}, 2, "f"); + f.define_extern("test", {}, {Int(32), Float(32)}, {x, y}); + f.realize({100, 100}); + printf("Success!\n"); + return 0; +} diff --git a/test/error/func_tuple_dim_mismatch.cpp b/test/error/func_tuple_dim_mismatch.cpp new file mode 100644 index 000000000000..79f97217d3b4 --- /dev/null +++ b/test/error/func_tuple_dim_mismatch.cpp @@ -0,0 +1,17 @@ +#include "Halide.h" +#include + +using namespace Halide; +using namespace Halide::Internal; + +int main(int argc, char **argv) { + Var x("x"), y("y"); + Func f({Int(32), Float(32)}, 1, "f"); + + f(x, y) = {cast(0), cast(0)}; + + f.realize({100, 100}); + + printf("Success!\n"); + return 0; +} diff --git a/test/error/func_tuple_types_mismatch.cpp b/test/error/func_tuple_types_mismatch.cpp new file mode 100644 index 000000000000..a04eaf45ce71 --- /dev/null +++ b/test/error/func_tuple_types_mismatch.cpp @@ -0,0 +1,17 @@ +#include "Halide.h" +#include + +using namespace Halide; +using namespace Halide::Internal; + +int main(int argc, char **argv) { + Var x("x"), y("y"); + Func f({UInt(8), Float(64)}, 2, "f"); + + f(x, y) = {cast(0), cast(0)}; + + f.realize({100, 100}); + + printf("Success!\n"); + return 0; +} diff --git a/test/error/func_tuple_update_types_mismatch.cpp b/test/error/func_tuple_update_types_mismatch.cpp new file mode 100644 index 000000000000..46577ce6ab66 --- /dev/null +++ b/test/error/func_tuple_update_types_mismatch.cpp @@ -0,0 +1,18 @@ +#include "Halide.h" +#include + +using namespace Halide; +using namespace Halide::Internal; + +int main(int argc, char **argv) { + Var x("x"), y("y"); + Func f({UInt(8), Float(64)}, 2, "f"); + + f(x, y) = {cast(0), cast(0)}; + f(x, y) += {cast(0), cast(0)}; + + f.realize({100, 100}); + + printf("Success!\n"); + return 0; +}