Infer input/output buffer constraints from schedule directives

apps/fft/main.cpp

@@ -153,6 +153,10 @@ int main(int argc, char **argv) {
     // locality.
     c2c_in(x, y, rep) = {re_in(x, y), im_in(x, y)};
     Func bench_c2c = fft2d_c2c(c2c_in, W, H, -1, target, fwd_desc);
+    // We must disable inference of buffer constraints for this Func,
+    // since we are going to use unorthodox strides later on. (We can
+    // still set buffer constraints manually.)
+    bench_c2c.infer_buffer_constraints(false);
     bench_c2c.compile_to_lowered_stmt(output_dir + "c2c.html", bench_c2c.infer_arguments(), HTML);
     Realization R_c2c = bench_c2c.realize(W, H, reps, target);
     // Write all reps to the same place in memory. This means the
@@ -184,6 +188,7 @@ int main(int argc, char **argv) {
     // All reps read from the same input. See notes on c2c_in.
     r2c_in(x, y, rep) = re_in(x, y);
     Func bench_r2c = fft2d_r2c(r2c_in, W, H, target, fwd_desc);
+    bench_r2c.infer_buffer_constraints(false);
     bench_r2c.compile_to_lowered_stmt(output_dir + "r2c.html", bench_r2c.infer_arguments(), HTML);
     Realization R_r2c = bench_r2c.realize(W, H/2 + 1, reps, target);
     // Write all reps to the same place in memory. See notes on R_c2c.
@@ -210,6 +215,7 @@ int main(int argc, char **argv) {
     // All reps read from the same input. See notes on c2c_in.
     c2r_in(x, y, rep) = {re_in(x, y), im_in(x, y)};
     Func bench_c2r = fft2d_c2r(c2r_in, W, H, target, inv_desc);
+    bench_c2r.infer_buffer_constraints(false);
     bench_c2r.compile_to_lowered_stmt(output_dir + "c2r.html", bench_c2r.infer_arguments(), HTML);
     Realization R_c2r = bench_c2r.realize(W, H, reps, target);
     // Write all reps to the same place in memory. See notes on R_c2c.

python_bindings/correctness/basics.py

@@ -28,7 +28,7 @@ def test_runtime_error():
     try:
         f.realize(buf)
     except RuntimeError as e:
-        assert 'do not cover required region' in str(e)
+        assert 'Constraint violated' in str(e)
     else:
         assert False, 'Did not see expected exception!'
 

python_bindings/correctness/buffer.py

@@ -173,6 +173,51 @@ def test_interleaved_ndarray():
     assert b.dim(2).extent() == c
     assert b.dim(2).stride() == 1
 
+def test_reorder():
+    W = 7
+    H = 5
+    C = 3
+    Z = 2
+
+    a = hl.Buffer(type = hl.UInt(8), sizes = [W, H, C], storage_order = [2, 0, 1])
+    assert a.dim(0).extent() == W
+    assert a.dim(1).extent() == H
+    assert a.dim(2).extent() == C
+    assert a.dim(2).stride() == 1
+    assert a.dim(0).stride() == C
+    assert a.dim(1).stride() == W * C
+
+    b = hl.Buffer(hl.UInt(8), [W, H, C, Z], [2, 3, 0, 1])
+    assert b.dim(0).extent() == W
+    assert b.dim(1).extent() == H
+    assert b.dim(2).extent() == C
+    assert b.dim(3).extent() == Z
+    assert b.dim(2).stride() == 1
+    assert b.dim(3).stride() == C
+    assert b.dim(0).stride() == C * Z
+    assert b.dim(1).stride() == W * C * Z
+
+    b2 = hl.Buffer(hl.UInt(8), [C, Z, W, H])
+    assert b.dim(0).extent() == b2.dim(2).extent()
+    assert b.dim(1).extent() == b2.dim(3).extent()
+    assert b.dim(2).extent() == b2.dim(0).extent()
+    assert b.dim(3).extent() == b2.dim(1).extent()
+    assert b.dim(0).stride() == b2.dim(2).stride()
+    assert b.dim(1).stride() == b2.dim(3).stride()
+    assert b.dim(2).stride() == b2.dim(0).stride()
+    assert b.dim(3).stride() == b2.dim(1).stride()
+
+    b2.transpose([2, 3, 0, 1])
+    assert b.dim(0).extent() == b2.dim(0).extent()
+    assert b.dim(1).extent() == b2.dim(1).extent()
+    assert b.dim(2).extent() == b2.dim(2).extent()
+    assert b.dim(3).extent() == b2.dim(3).extent()
+    assert b.dim(0).stride() == b2.dim(0).stride()
+    assert b.dim(1).stride() == b2.dim(1).stride()
+    assert b.dim(2).stride() == b2.dim(2).stride()
+    assert b.dim(3).stride() == b2.dim(3).stride()
+
+
 if __name__ == "__main__":
     test_make_interleaved()
     test_interleaved_ndarray()
@@ -182,4 +227,5 @@ def test_interleaved_ndarray():
     test_fill_all_equal()
     test_bufferinfo_sharing()
     test_float16()
+    test_reorder()
 

python_bindings/src/PyBuffer.cpp

@@ -304,6 +304,10 @@ void define_buffer(py::module &m) {
             return Buffer<>(type, sizes, name);
         }), py::arg("type"), py::arg("sizes"), py::arg("name") = "")
 
+        .def(py::init([](Type type, const std::vector<int> &sizes, const std::vector<int> &storage_order, const std::string &name) -> Buffer<> {
+            return Buffer<>(type, sizes, storage_order, name);
+        }), py::arg("type"), py::arg("sizes"), py::arg("storage_order"), py::arg("name") = "")
+
         // Note that this exists solely to allow you to create a Buffer with a null host ptr;
         // this is necessary for some bounds-query operations (e.g. Func::infer_input_bounds).
         .def_static("make_bounds_query", [](Type type, const std::vector<int> &sizes, const std::string &name) -> Buffer<> {
@@ -419,10 +423,17 @@ void define_buffer(py::module &m) {
             b.transpose(d1, d2);
         }, py::arg("d1"), py::arg("d2"))
 
-        // Present in Runtime::Buffer but not Buffer
-        // .def("transposed", [](Buffer<> &b, int d1, int d2) -> Buffer<> {
-        //     return b.transposed(d1, d2);
-        // }, py::arg("d1"), py::arg("d2"))
+        .def("transposed", [](Buffer<> &b, int d1, int d2) -> Buffer<> {
+            return b.transposed(d1, d2);
+        }, py::arg("d1"), py::arg("d2"))
+
+        .def("transpose", [](Buffer<> &b, const std::vector<int> &order) -> void {
+            b.transpose(order);
+        }, py::arg("order"))
+
+        .def("transposed", [](Buffer<> &b, const std::vector<int> &order) -> Buffer<> {
+            return b.transposed(order);
+        }, py::arg("order"))
 
         .def("dim", [](Buffer<> &b, int dimension) -> BufferDimension {
             return b.dim(dimension);

src/AddImageChecks.cpp

@@ -1,4 +1,5 @@
 #include "AddImageChecks.h"
+#include "IREquality.h"
 #include "IRVisitor.h"
 #include "Simplify.h"
 #include "Substitute.h"
@@ -189,6 +190,7 @@ Stmt add_image_checks(Stmt s,
         bool is_output_buffer = false;
         bool is_secondary_output_buffer = false;
         string buffer_name = name;
+        Parameter primary_output_buffer;
         for (Function f : outputs) {
             for (size_t i = 0; i < f.output_buffers().size(); i++) {
                 if (param.defined() &&
@@ -199,6 +201,7 @@ Stmt add_image_checks(Stmt s,
                     buffer_name = f.name();
                     if (i > 0) {
                         is_secondary_output_buffer = true;
+                        primary_output_buffer = f.output_buffers()[0];
                     }
                 }
             }
@@ -449,8 +452,14 @@ Stmt add_image_checks(Stmt s,
                 // constrained to match the first output.
 
                 if (param.defined()) {
-                    user_assert(!param.extent_constraint(i).defined() &&
-                                !param.min_constraint(i).defined())
+                    // Parameter::set_constraints_from_schedule() can set the min/extent
+                    // on secondary buffers, so don't complain if they are identical
+                    // to the primary.
+                    const auto constraint_ok = [](const Expr &secondary, const Expr &primary) -> bool {
+                        return !secondary.defined() || equal(secondary, primary);
+                    };
+                    user_assert(constraint_ok(param.extent_constraint(i), primary_output_buffer.extent_constraint(i)) &&
+                                constraint_ok(param.min_constraint(i), primary_output_buffer.min_constraint(i)))
                         << "Can't constrain the min or extent of an output buffer beyond the "
                         << "first. They are implicitly constrained to have the same min and extent "
                         << "as the first output buffer.\n";

src/Buffer.h

@@ -185,10 +185,21 @@ class Buffer {
                     const std::string &name = "") :
         Buffer(Runtime::Buffer<T>(t, sizes), name) {}
 
+    explicit Buffer(Type t,
+                    const std::vector<int> &sizes,
+                    const std::vector<int> &storage_order,
+                    const std::string &name = "") :
+        Buffer(Runtime::Buffer<T>(t, sizes, storage_order), name) {}
+
     explicit Buffer(const std::vector<int> &sizes,
                     const std::string &name = "") :
         Buffer(Runtime::Buffer<T>(sizes), name) {}
 
+    explicit Buffer(const std::vector<int> &sizes,
+                    const std::vector<int> &storage_order,
+                    const std::string &name = "") :
+        Buffer(Runtime::Buffer<T>(sizes, storage_order), name) {}
+
     template<typename Array, size_t N>
     explicit Buffer(Array (&vals)[N],
                     const std::string &name = "") :
@@ -395,6 +406,7 @@ class Buffer {
     HALIDE_BUFFER_FORWARD(translate)
     HALIDE_BUFFER_FORWARD_INITIALIZER_LIST(translate, std::vector<int>)
     HALIDE_BUFFER_FORWARD(transpose)
+    HALIDE_BUFFER_FORWARD(transposed)
     HALIDE_BUFFER_FORWARD(add_dimension)
     HALIDE_BUFFER_FORWARD(copy_to_host)
     HALIDE_BUFFER_FORWARD(copy_to_device)

src/Func.cpp

@@ -2480,6 +2480,12 @@ Func &Func::add_trace_tag(const std::string &trace_tag) {
     return *this;
 }
 
+Func &Func::infer_buffer_constraints(bool infer) {
+    invalidate_cache();
+    func.infer_buffer_constraints() = infer;
+    return *this;
+}
+
 void Func::debug_to_file(const string &filename) {
     invalidate_cache();
     func.debug_file() = filename;

src/Func.h

@@ -2301,6 +2301,12 @@ class Func {
      */
     Func &add_trace_tag(const std::string &trace_tag);
 
+    /** Specify whether constraints for the buffer underlying this function
+     * should be inferred from other scheduling directives (e.g. to determine
+     * storage layout, strides, etc). Defaults to true. It should be quite
+     * rare to need to disable this. */
+    Func &infer_buffer_constraints(bool infer);
+
     /** Get a handle on the internal halide function that this Func
      * represents. Useful if you want to do introspection on Halide
      * functions */

src/Function.cpp

@@ -91,6 +91,7 @@ struct FunctionContents {
     std::vector<string> trace_tags;
 
     bool frozen = false;
+    bool infer_buffer_constraints = true;
 
     void accept(IRVisitor *visitor) const {
         func_schedule.accept(visitor);
@@ -343,6 +344,7 @@ void Function::deep_copy(FunctionPtr copy, DeepCopyMap &copied_map) const {
     copy->trace_realizations = contents->trace_realizations;
     copy->trace_tags = contents->trace_tags;
     copy->frozen = contents->frozen;
+    copy->infer_buffer_constraints = contents->infer_buffer_constraints;
     copy->output_buffers = contents->output_buffers;
     copy->func_schedule = contents->func_schedule.deep_copy(copied_map);
 
@@ -796,6 +798,24 @@ const FuncSchedule &Function::schedule() const {
     return contents->func_schedule;
 }
 
+std::vector<int> Function::storage_order() const {
+    const std::vector<StorageDim> &storage_dims = contents->func_schedule.storage_dims();
+    const std::vector<std::string> &args = contents->args;
+    internal_assert(storage_dims.size() == args.size());
+
+    std::vector<int> storage_order;
+    for (const auto &s : storage_dims) {
+        for (size_t a = 0; a < args.size(); ++a) {
+            if (args[a] == s.var) {
+                storage_order.push_back((int) a);
+                break;
+            }
+        }
+    }
+    internal_assert(storage_order.size() == args.size());
+    return storage_order;
+}
+
 const std::vector<Parameter> &Function::output_buffers() const {
     return contents->output_buffers;
 }
@@ -896,6 +916,14 @@ std::string &Function::debug_file() {
     return contents->debug_file;
 }
 
+bool Function::infer_buffer_constraints() const {
+    return contents->infer_buffer_constraints;
+}
+
+bool &Function::infer_buffer_constraints() {
+    return contents->infer_buffer_constraints;
+}
+
 void Function::trace_loads() {
     contents->trace_loads = true;
 }

src/Function.h

@@ -186,6 +186,11 @@ class Function {
     /** Get a const handle to the function-specific schedule for inspecting it. */
     const FuncSchedule &schedule() const;
 
+    /** Return a vector containing the order of storage for each argument,
+     * based on the scheduling directives. Each entry is an index into the args array.
+     */
+    std::vector<int> storage_order() const;
+
     /** Get a handle on the output buffer used for setting constraints
      * on it. */
     const std::vector<Parameter> &output_buffers() const;
@@ -267,6 +272,14 @@ class Function {
     /** Get a handle to the debug filename. */
     std::string &debug_file();
 
+    /** If true, attempt to infer constraints on the buffer shape based
+     * on the scheduling directives. (Default = true; it is very rare to
+     * need to disable this.) */
+    // @{
+    bool infer_buffer_constraints() const;
+    bool &infer_buffer_constraints();
+    // @}
+
     /** Use an an extern argument to another function. */
     operator ExternFuncArgument() const {
         return ExternFuncArgument(contents);

src/Generator.cpp

@@ -152,6 +152,7 @@ Func make_param_func(const Parameter &p, const std::string &name) {
         }
         f(args) = Internal::Call::make(p, args_expr);
     }
+    f.function().schedule().emit_inliner_warnings() = false;
     return f;
 }
 
@@ -1165,6 +1166,7 @@ Func GeneratorBase::get_output(const std::string &n) {
     user_assert(!output->is_array() && output->funcs().size() == 1) << "Output " << n << " must be accessed via get_array_output()\n";
     Func f = output->funcs().at(0);
     user_assert(f.defined()) << "Output " << n << " was not defined.\n";
+    f.function().schedule().emit_inliner_warnings() = false;
     return f;
 }
 
@@ -1175,6 +1177,7 @@ std::vector<Func> GeneratorBase::get_array_output(const std::string &n) {
     (void) output->array_size();
     for (const auto &f : output->funcs()) {
         user_assert(f.defined()) << "Output " << n << " was not fully defined.\n";
+        f.function().schedule().emit_inliner_warnings() = false;
     }
     return output->funcs();
 }
@@ -1390,6 +1393,14 @@ Module GeneratorBase::build_module(const std::string &function_name,
     ParamInfo &pi = param_info();
     std::vector<Argument> filter_arguments;
     for (auto input : pi.filter_inputs) {
+        if (input->kind() == IOKind::Function) {
+            internal_assert(input->parameters_.size() == input->funcs_.size());
+            for (size_t i = 0; i < input->parameters_.size(); ++i) {
+                Function f = input->funcs_.at(i).function();
+                // Transfer the constraints from the Function to the Parameter
+                input->parameters_.at(i).set_constraints_from_schedule(f);
+            }
+        }
         for (const auto &p : input->parameters_) {
             filter_arguments.push_back(to_argument(p, p.is_buffer() ? Expr() : input->get_def_expr()));
         }
@@ -1695,6 +1706,7 @@ void GeneratorInputBase::set_inputs(const std::vector<StubInput> &inputs) {
         if (kind() == IOKind::Function) {
             auto f = in.func();
             user_assert(f.defined()) << "The input for " << name() << " is an undefined Func. Please define it.\n";
+            f.function().schedule().emit_inliner_warnings() = false;
             check_matching_types(f.output_types());
             check_matching_dims(f.dimensions());
             funcs_.push_back(f);
@@ -1998,7 +2010,7 @@ void generator_test() {
         static_assert(std::is_same<decltype(tester_instance.expr_array_input[0]), const Expr &>::value, "type mismatch");
         static_assert(std::is_same<decltype(tester_instance.expr_array_output[0]), const Expr &>::value, "type mismatch");
 
-        static_assert(std::is_same<decltype(tester_instance.func_array_input[0]), const Func &>::value, "type mismatch");
+        static_assert(std::is_same<decltype(tester_instance.func_array_input[0]), Func>::value, "type mismatch");
         static_assert(std::is_same<decltype(tester_instance.func_array_output[0]), Func &>::value, "type mismatch");
 
         static_assert(std::is_same<decltype(tester_instance.buffer_array_input[0]), ImageParam>::value, "type mismatch");