Add a new, alternate JIT-call convention (#6777)

* Prototype of revised JIT-call convention

Experiment to try out a way to call JIT code in C++ using the same calling conventions as AOT code. Very much experimental.

* Update Pipeline.h

* Add Python support for `compile_to_callable` + make empty_ucon static

* Update PyCallable.cpp

* Update buffer.py

* wip

* Update callable.py

* WIP

* Update custom_allocator.cpp

* Update Callable.cpp

* Add Generator support for Callables

* Update Generator.cpp

* Update PyPipeline.cpp

* Fixes

* Update callable.cpp

* Update CMakeLists.txt

* create_callable_from_generator

* More cleanup

* Update Generator.cpp

* Fix Python bounds inference

* Add Python wrapper for create_callable_from_generator() + Add kwarg support for Callable

* Add set_generatorparam_values() + usage

* Fix auto_schedule/machine_params parsing

The recent refactoring that added `execute_generator` accidentally nuked setting these two GeneratorParams. Oops. Fixed.

* Move the type-checking code into a constexpr code

* Update Callable.h

* clang-tidy

* CLANG-TIDY

* Add `make_std_function`, + more general cleanup

* Update example_jittest.cpp

* Update Callable.h

* Update Callable.h

* More tweaking, smaller CallCheckInfo

* Still more cleanup

* make_std_function now does Buffer type/dim checking where possible

* Add tests for calling `AbstractGenreator::compile_to_callable()` directly

* enable exports

* Various fixes

* Improve fill_slot for Halide::Buffer

* kill report_if_error

* Update callable_bad_arguments.cpp

* Update Pipeline.cpp

* Revise error handling

* Update Callable.cpp

* Update callable.py

* Update callable_generator.cpp

* Update callable.py

* HALIDE_MUST_USE_RESULT -> HALIDE_FUNCTION_ATTRS for Callable

Makefile

@@ -417,6 +417,7 @@ SOURCE_FILES = \
   BoundsInference.cpp \
   BoundSmallAllocations.cpp \
   Buffer.cpp \
+  Callable.cpp \
   CanonicalizeGPUVars.cpp \
   Closure.cpp \
   ClampUnsafeAccesses.cpp \
@@ -594,6 +595,7 @@ HEADER_FILES = \
   BoundsInference.h \
   BoundSmallAllocations.h \
   Buffer.h \
+  Callable.h \
   CanonicalizeGPUVars.h \
   ClampUnsafeAccesses.h \
   Closure.h \

python_bindings/correctness/CMakeLists.txt

@@ -11,6 +11,7 @@ set(TESTS
     bit_test.py
     boundary_conditions.py
     buffer.py
+    callable.py
     compile_to.py
     division.py
     extern.py

python_bindings/correctness/buffer.py

@@ -250,7 +250,7 @@ def test_overflow():
     try:
         hl.Buffer(size_over_intmax)
     except ValueError as e:
-        assert 'Out of range arguments to make_dim_vec.' in str(e)
+        assert 'Out of range dimensions in buffer conversion' in str(e)
 
 def test_buffer_to_str():
     b = hl.Buffer()

python_bindings/correctness/callable.py

@@ -0,0 +1,177 @@
+import halide as hl
+import numpy as np
+
+import simple_pystub  # Needed for create_callable_from_generator("simple") to work
+
+def test_callable():
+    p_int16 = hl.Param(hl.Int(16), 42)
+    p_float = hl.Param(hl.Float(32), 1.0)
+    p_img = hl.ImageParam(hl.UInt(8), 2)
+
+    x = hl.Var('x')
+    y = hl.Var('y')
+    f = hl.Func('f')
+
+    f[x, y] = p_img[x, y] + hl.u8(p_int16 / p_float)
+
+    in1 = hl.Buffer(hl.UInt(8), [10, 10])
+    in2 = hl.Buffer(hl.UInt(8), [10, 10])
+
+    for i in range(10):
+        for j in range(10):
+            in1[i, j] = i + j * 10
+            in2[i, j] = i * 10 + j
+
+    c = f.compile_to_callable([p_img, p_int16, p_float]);
+
+    out1 = hl.Buffer(hl.UInt(8), [10, 10])
+    c(in1, 42, 1.0, out1)
+
+    out2 = hl.Buffer(hl.UInt(8), [10, 10])
+    c(in2, 22, 2.0, out2)
+
+    out3 = hl.Buffer(hl.UInt(8), [10, 10])
+    c(in1, 12, 1.0, out3)
+
+    out4 = hl.Buffer(hl.UInt(8), [10, 10])
+    c(in2, 16, 1.0, out4)
+
+    for i in range(10):
+        for j in range(10):
+            assert out1[i, j] == i + j * 10 + 42
+            assert out2[i, j] == i * 10 + j + 11
+            assert out3[i, j] == i + j * 10 + 12
+            assert out4[i, j] == i * 10 + j + 16
+
+    # Test bounds inference. Note that in Python there
+    # isn't a "natural" way to create a buffer with a null host ptr
+    # so we use this specific API for the purpose.
+    in_bounds = hl.Buffer.make_bounds_query(hl.UInt(8), [1, 1])
+    out_bounds = hl.Buffer.make_bounds_query(hl.UInt(8), [20, 20])
+    c(in_bounds, 42, 1.0, out_bounds)
+
+    assert in_bounds.defined()
+    assert in_bounds.dim(0).extent() == 20
+    assert in_bounds.dim(1).extent() == 20
+    assert in1.dim(0).extent() == 10
+    assert in1.dim(1).extent() == 10
+
+def test_simple():
+    x, y = hl.Var(), hl.Var()
+    target = hl.get_jit_target_from_environment()
+
+    b_in = hl.Buffer(hl.UInt(8), [2, 2])
+    b_in.fill(123)
+
+    # All inputs to a Callable must be fully realized, so any Func inputs
+    # that the Generator has implicitly become Buffer inputs of the same type
+    # and dimensionality.
+    f_in = hl.Buffer(hl.Int(32), [2, 2])
+    for xx in range(2):
+        for yy in range(2):
+                f_in[xx, yy] = xx + yy
+
+    float_in = 3.5
+
+    b_out = hl.Buffer(hl.Float(32), [2, 2])
+
+    def _check(offset = 0):
+        assert b_out[0, 0] == float_in + 0 + offset + 123
+        assert b_out[0, 1] == float_in + 1 + offset + 123
+        assert b_out[1, 0] == float_in + 1 + offset + 123
+        assert b_out[1, 1] == float_in + 2 + offset + 123
+
+    gp = {"func_input.type": "int32"}
+    simple = hl.create_callable_from_generator(target, "simple", gp)
+
+    # ----------- Positional arguments
+    simple(b_in, f_in, float_in, b_out)
+    _check()
+
+    # ----------- Keyword arguments
+    # Natural order
+    simple(buffer_input=b_in, func_input=f_in, float_arg=float_in, simple_output=b_out)
+    _check()
+
+    # Weird order
+    simple(float_arg=float_in, simple_output=b_out, buffer_input=b_in, func_input=f_in)
+    _check()
+
+    # ----------- Positional + Keywords
+
+    # Natural order
+    simple(b_in, func_input=f_in, simple_output=b_out, float_arg=float_in)
+    _check()
+
+    # Weird order
+    simple(b_in, f_in, float_in, simple_output=b_out)
+    _check()
+
+    # ----------- Above set again, w/ additional GeneratorParam mixed in
+    k = 42
+
+    gp = {"func_input.type": "int32", "offset": str(k)}
+    simple_42 = hl.create_callable_from_generator(target, "simple", gp)
+    simple_42(b_in, f_in, float_in, b_out)
+    _check(k)
+
+    # ----------- Test various failure modes
+    try:
+        # too many positional args
+        simple(b_in, f_in, float_in, 4, b_out)
+    except hl.HalideError as e:
+        assert 'Expected at most 4 positional arguments, but saw 5.' in str(e)
+    else:
+        assert False, 'Did not see expected exception!'
+
+    try:
+        # too few positional args
+        simple(b_in, f_in)
+    except hl.HalideError as e:
+        assert 'Expected exactly 4 positional arguments, but saw 2.' in str(e)
+    else:
+        assert False, 'Did not see expected exception!'
+
+    try:
+        # Inputs that can't be converted to what the receiver needs (positional)
+        simple(hl.f32(3.141592), "happy", k, b_out)
+    except hl.HalideError as e:
+        assert 'is not an instance of' in str(e)
+    else:
+        assert False, 'Did not see expected exception!'
+
+    try:
+        # Inputs that can't be converted to what the receiver needs (named)
+        simple(b_in, f_in, float_in, simple_output="bogus")
+    except hl.HalideError as e:
+        assert 'is not an instance of' in str(e)
+    else:
+        assert False, 'Did not see expected exception!'
+
+    try:
+        # Bad keyword argument
+        simple(buffer_input=b_in, float_arg=float_in, simple_output=b_out, funk_input=f_in)
+    except hl.HalideError as e:
+        assert "Unknown argument 'funk_input' specified via keyword." in str(e)
+    else:
+        assert False, 'Did not see expected exception!'
+
+    try:
+        # too few keyword args
+        simple(float_arg=float_in, simple_output=b_out, func_input=f_in)
+    except hl.HalideError as e:
+        assert 'Argument buffer_input was not specified by either positional or keyword argument.' in str(e)
+    else:
+        assert False, 'Did not see expected exception!'
+
+    try:
+        # Arg specified by pos + kw
+        simple(b_in, buffer_input=b_in, func_input=f_in, float_arg=float_in, simple_output=b_out)
+    except hl.HalideError as e:
+        assert 'Argument buffer_input specified multiple times.' in str(e)
+    else:
+        assert False, 'Did not see expected exception!'
+
+if __name__ == "__main__":
+    test_callable()
+    test_simple()

python_bindings/src/CMakeLists.txt

@@ -2,6 +2,7 @@ set(SOURCES
     PyArgument.cpp
     PyBoundaryConditions.cpp
     PyBuffer.cpp
+    PyCallable.cpp
     PyConciseCasts.cpp
     PyDerivative.cpp
     PyEnums.cpp

python_bindings/src/PyBuffer.cpp

@@ -143,6 +143,8 @@ std::string type_to_format_descriptor(const Type &type) {
     return std::string();
 }
 
+}  // namespace
+
 Type format_descriptor_to_type(const std::string &fd) {
 
 #define HANDLE_BUFFER_TYPE(TYPE) \
@@ -173,6 +175,8 @@ Type format_descriptor_to_type(const std::string &fd) {
     return Type();
 }
 
+namespace {
+
 py::object buffer_getitem_operator(Buffer<> &buf, const std::vector<int> &pos) {
     if ((size_t)pos.size() != (size_t)buf.dimensions()) {
         throw py::value_error("Incorrect number of dimensions.");
@@ -236,26 +240,8 @@ py::object buffer_setitem_operator(Buffer<> &buf, const std::vector<int> &pos, c
 class PyBuffer : public Buffer<> {
     py::buffer_info info;
 
-    static std::vector<halide_dimension_t> make_dim_vec(const py::buffer_info &info) {
-        const Type t = format_descriptor_to_type(info.format);
-        std::vector<halide_dimension_t> dims;
-        dims.reserve(info.ndim);
-        for (int i = 0; i < info.ndim; i++) {
-            if (INT_MAX < info.shape[i] || INT_MAX < (info.strides[i] / t.bytes())) {
-                throw py::value_error("Out of range arguments to make_dim_vec.");
-            }
-            dims.emplace_back(0, (int32_t)info.shape[i], (int32_t)(info.strides[i] / t.bytes()));
-        }
-        return dims;
-    }
-
     PyBuffer(py::buffer_info &&info, const std::string &name)
-        : Buffer<>(
-              format_descriptor_to_type(info.format),
-              info.ptr,
-              (int)info.ndim,
-              make_dim_vec(info).data(),
-              name),
+        : Buffer<>(pybufferinfo_to_halidebuffer(info), name),
           info(std::move(info)) {
     }
 

python_bindings/src/PyBuffer.h

@@ -8,6 +8,31 @@ namespace PythonBindings {
 
 void define_buffer(py::module &m);
 
+Type format_descriptor_to_type(const std::string &fd);
+
+template<typename T = void,
+         int Dims = AnyDims,
+         int InClassDimStorage = (Dims == AnyDims ? 4 : std::max(Dims, 1))>
+Halide::Runtime::Buffer<T, Dims, InClassDimStorage> pybufferinfo_to_halidebuffer(const py::buffer_info &info) {
+    const Type t = format_descriptor_to_type(info.format);
+    halide_dimension_t *dims = (halide_dimension_t *)alloca(info.ndim * sizeof(halide_dimension_t));
+    _halide_user_assert(dims);
+    for (int i = 0; i < info.ndim; i++) {
+        if (INT_MAX < info.shape[i] || INT_MAX < (info.strides[i] / t.bytes())) {
+            throw py::value_error("Out of range dimensions in buffer conversion.");
+        }
+        dims[i] = {0, (int32_t)info.shape[i], (int32_t)(info.strides[i] / t.bytes())};
+    }
+    return Halide::Runtime::Buffer<T, Dims, InClassDimStorage>(t, info.ptr, (int)info.ndim, dims);
+}
+
+template<typename T = void,
+         int Dims = AnyDims,
+         int InClassDimStorage = (Dims == AnyDims ? 4 : std::max(Dims, 1))>
+Halide::Runtime::Buffer<T, Dims, InClassDimStorage> pybuffer_to_halidebuffer(const py::buffer &pyb, bool writable) {
+    return pybufferinfo_to_halidebuffer(pyb.request(writable));
+}
+
 }  // namespace PythonBindings
 }  // namespace Halide