Allow PyPipeline and PyFunc to realize() scalar buffers

python_bindings/correctness/complexstub_generator.cpp

@@ -35,6 +35,7 @@ class ComplexStub : public Halide::Generator<ComplexStub> {
  Output<Buffer<float, 3>> typed_buffer_output{"typed_buffer_output"};
  Output<Buffer<void, 3>> untyped_buffer_output{"untyped_buffer_output"};
  Output<Buffer<uint8_t, 3>> static_compiled_buffer_output{"static_compiled_buffer_output"};
+ Output<float> scalar_output{"scalar_output"};
 
  void configure() {
  // Pointers returned by add_input() are managed by the Generator;
@@ -72,6 +73,8 @@ class ComplexStub : public Halide::Generator<ComplexStub> {
  static_compiled_buffer_output = static_compiled_buffer;
 
  (*extra_func_output)(x, y) = cast<double>((*extra_func_input)(x, y, 0) + 1);
+
+ scalar_output() = float_arg + int_arg;
  }
 
  void schedule() {

python_bindings/correctness/pystub.py

@@ -184,6 +184,7 @@ def test_complexstub():
  typed_buffer_output,
  untyped_buffer_output,
  static_compiled_buffer_output,
+ scalar_output,
  extra_func_output) = r
 
  b = simple_output.realize([32, 32, 3], target)
@@ -249,6 +250,10 @@ def test_complexstub():
  actual = b[x, y, c]
  assert expected == actual, "Expected %s Actual %s" % (expected, actual)
 
+ b = scalar_output.realize([], target)
+ assert b.type() == hl.Float(32)
+ assert b[()] == 34.25
+
  b = extra_func_output.realize([32, 32], target)
  assert b.type() == hl.Float(64)
  for x in range(32):

python_bindings/src/PyFunc.cpp

@@ -122,15 +122,12 @@ void define_func(py::module &m) {
  },
  py::arg("dst"), py::arg("target") = Target())
 
- // This will actually allow a list-of-buffers as well as a tuple-of-buffers, but that's OK.
- .def(
- "realize",
- [](Func &f, std::vector<Buffer<>> buffers, const Target &t) -> void {
- py::gil_scoped_release release;
- f.realize(Realization(buffers), t);
- },
- py::arg("dst"), py::arg("target") = Target())
-
+ // It's important to have this overload of realize() go first:
+ // passing an empty list [] is ambiguous in Python, and could match to
+ // either list-of-sizes or list-of-buffers... but the former is useful
+ // (it allows realizing a 0-dimensional/scalar buffer) and the former is
+ // not (it will always assert-fail). Putting this one first allows it to
+ // be the first one chosen by the bindings in this case.
  .def(
  "realize",
  [](Func &f, const std::vector<int32_t> &sizes, const Target &target) -> py::object {
@@ -143,6 +140,15 @@ void define_func(py::module &m) {
  },
  py::arg("sizes") = std::vector<int32_t>{}, py::arg("target") = Target())
 
+ // This will actually allow a list-of-buffers as well as a tuple-of-buffers, but that's OK.
+ .def(
+ "realize",
+ [](Func &f, std::vector<Buffer<>> buffers, const Target &t) -> void {
+ py::gil_scoped_release release;
+ f.realize(Realization(buffers), t);
+ },
+ py::arg("dst"), py::arg("target") = Target())
+
  .def("defined", &Func::defined)
  .def("name", &Func::name)
  .def("dimensions", &Func::dimensions)

python_bindings/src/PyPipeline.cpp

@@ -101,14 +101,12 @@ void define_pipeline(py::module &m) {
  },
  py::arg("dst"), py::arg("target") = Target())
 
- // This will actually allow a list-of-buffers as well as a tuple-of-buffers, but that's OK.
- .def(
- "realize", [](Pipeline &p, std::vector<Buffer<>> buffers, const Target &t) -> void {
- py::gil_scoped_release release;
- p.realize(Realization(buffers), t);
- },
- py::arg("dst"), py::arg("target") = Target())
-
+ // It's important to have this overload of realize() go first:
+ // passing an empty list [] is ambiguous in Python, and could match to
+ // either list-of-sizes or list-of-buffers... but the former is useful
+ // (it allows realizing a 0-dimensional/scalar buffer) and the former is
+ // not (it will always assert-fail). Putting this one first allows it to
+ // be the first one chosen by the bindings in this case.
  .def(
  "realize", [](Pipeline &p, std::vector<int32_t> sizes, const Target &target) -> py::object {
  std::optional<Realization> r;
@@ -120,6 +118,14 @@ void define_pipeline(py::module &m) {
  },
  py::arg("sizes") = std::vector<int32_t>{}, py::arg("target") = Target())
 
+ // This will actually allow a list-of-buffers as well as a tuple-of-buffers, but that's OK.
+ .def(
+ "realize", [](Pipeline &p, std::vector<Buffer<>> buffers, const Target &t) -> void {
+ py::gil_scoped_release release;
+ p.realize(Realization(buffers), t);
+ },
+ py::arg("dst"), py::arg("target") = Target())
+
  .def(
  "infer_input_bounds", [](Pipeline &p, const py::object &dst, const Target &target) -> void {
  // dst could be Buffer<>, vector<Buffer>, or vector<int>