Narrowing/widening the Func along a dimension

As discussed, this implements an utility to make it easier
to write code that wants to essentially operate either on
smaller chunks than the whole type, or operate on several
consecutive elements as a large element.

The bigger picture is that it is desirable to perform load
widening in such situations, and while this doesn't do that,
at least having a common interface should be a step in that direction.

I've took liberty to add/expose some QoL variable-bit-lenth
operations in IROperator, while there.

I believe, this has sufficient test coverage now.
While implmeneting, i stumbled&fixed
#6782

Refs. #6756

python_bindings/src/CMakeLists.txt

@@ -10,6 +10,7 @@ set(SOURCES
     PyExternFuncArgument.cpp
     PyFunc.cpp
     PyFuncRef.cpp
+    PyFuncTypeChanging.cpp
     PyHalide.cpp
     PyImageParam.cpp
     PyInlineReductions.cpp

python_bindings/src/PyFuncTypeChanging.cpp

@@ -0,0 +1,66 @@
+#include "PyFuncTypeChanging.h"
+
+namespace Halide {
+namespace PythonBindings {
+
+namespace {
+
+inline Func to_func(const Buffer<> &b) {
+    return lambda(_, b(_));
+}
+
+}  // namespace
+
+void define_func_type_changing(py::module &m) {
+    using namespace FuncTypeChanging;
+
+    py::module bc = m.def_submodule("FuncTypeChanging");
+
+    py::enum_<ChunkOrder>(bc, "ArgumentKind")
+        .value("LowestFirst", ChunkOrder::LowestFirst)
+        .value("HighestFirst", ChunkOrder::HighestFirst)
+        .value("Default", ChunkOrder::Default);
+
+    bc.def(
+        "change_type",
+        [](const ImageParam &im, const Type &dst_type, const Var &dim,
+           const std::string &name, ChunkOrder chunk_order) -> Func {
+            return change_type(im, dst_type, dim, name, chunk_order);
+        },
+        py::arg("f"), py::arg("dst_type"), py::arg("dim"), py::arg("name"),
+        py::arg("chunk_order"));
+
+    bc.def(
+        "change_type",
+        [](const Buffer<> &b, const Type &dst_type, const Var &dim,
+           const std::string &name, ChunkOrder chunk_order) -> Func {
+            return change_type(b, dst_type, dim, name, chunk_order);
+        },
+        py::arg("f"), py::arg("dst_type"), py::arg("dim"), py::arg("name"),
+        py::arg("chunk_order"));
+
+    bc.def(
+        "change_type",
+        [](const py::object &target, const Type &dst_type, const Var &dim,
+           const std::string &name, ChunkOrder chunk_order) -> Func {
+            try {
+                return change_type(target.cast<Func>(), dst_type, dim, name,
+                                   chunk_order);
+            } catch (...) {
+                // fall thru
+            }
+            try {
+                return change_type(to_func(target.cast<Buffer<>>()), dst_type,
+                                   dim, name, chunk_order);
+            } catch (...) {
+                // fall thru
+            }
+            throw py::value_error("Invalid arguments to change_type");
+            return Func();
+        },
+        py::arg("f"), py::arg("dst_type"), py::arg("dim"), py::arg("name"),
+        py::arg("chunk_order"));
+}
+
+}  // namespace PythonBindings
+}  // namespace Halide

python_bindings/src/PyFuncTypeChanging.h

@@ -0,0 +1,14 @@
+#ifndef HALIDE_PYTHON_BINDINGS_PYFUNCTYPECHANGING_H
+#define HALIDE_PYTHON_BINDINGS_PYFUNCTYPECHANGING_H
+
+#include "PyHalide.h"
+
+namespace Halide {
+namespace PythonBindings {
+
+void define_func_type_changing(py::module &m);
+
+}  // namespace PythonBindings
+}  // namespace Halide
+
+#endif  // HALIDE_PYTHON_BINDINGS_PYFUNCTYPECHANGING_H

python_bindings/src/PyHalide.cpp

@@ -10,6 +10,7 @@
 #include "PyExpr.h"
 #include "PyExternFuncArgument.h"
 #include "PyFunc.h"
+#include "PyFuncTypeChanging.h"
 #include "PyIROperator.h"
 #include "PyImageParam.h"
 #include "PyInlineReductions.h"
@@ -46,6 +47,7 @@ PYBIND11_MODULE(HALIDE_PYBIND_MODULE_NAME, m) {
     define_tuple(m);
     define_argument(m);
     define_boundary_conditions(m);
+    define_func_type_changing(m);
     define_buffer(m);
     define_concise_casts(m);
     define_error(m);

python_bindings/src/PyIROperator.cpp

@@ -167,6 +167,10 @@ void define_operators(py::module &m) {
     m.def("popcount", &popcount);
     m.def("count_leading_zeros", &count_leading_zeros);
     m.def("count_trailing_zeros", &count_trailing_zeros);
+    m.def("extract_high_bits", &extract_high_bits);
+    m.def("variable_length_extend", &variable_length_extend);
+    m.def("extract_bits", &extract_bits);
+    m.def("extract_low_bits", &extract_low_bits);
     m.def("div_round_to_zero", &div_round_to_zero);
     m.def("mod_round_to_zero", &mod_round_to_zero);
     m.def("random_float", (Expr(*)()) & random_float);

src/BoundaryConditions.h

@@ -9,6 +9,7 @@
 
 #include "Expr.h"
 #include "Func.h"
+#include "FuncExtras.h"
 #include "Lambda.h"
 
 namespace Halide {
@@ -62,15 +63,6 @@ inline HALIDE_NO_USER_CODE_INLINE void collect_region(Region &collected_args,
     collect_region(collected_args, std::forward<Args>(args)...);
 }
 
-inline const Func &func_like_to_func(const Func &func) {
-    return func;
-}
-
-template<typename T>
-inline HALIDE_NO_USER_CODE_INLINE Func func_like_to_func(const T &func_like) {
-    return lambda(_, func_like(_));
-}
-
 }  // namespace Internal
 
 /** Impose a boundary condition such that a given expression is returned
@@ -99,12 +91,12 @@ Func constant_exterior(const Func &source, const Expr &value,
 
 template<typename T>
 HALIDE_NO_USER_CODE_INLINE Func constant_exterior(const T &func_like, const Tuple &value, const Region &bounds) {
-    return constant_exterior(Internal::func_like_to_func(func_like), value, bounds);
+    return constant_exterior(::Halide::Internal::func_like_to_func(func_like), value, bounds);
 }
 
 template<typename T>
 HALIDE_NO_USER_CODE_INLINE Func constant_exterior(const T &func_like, const Expr &value, const Region &bounds) {
-    return constant_exterior(Internal::func_like_to_func(func_like), value, bounds);
+    return constant_exterior(::Halide::Internal::func_like_to_func(func_like), value, bounds);
 }
 
 template<typename T>
@@ -114,7 +106,7 @@ HALIDE_NO_USER_CODE_INLINE Func constant_exterior(const T &func_like, const Tupl
         object_bounds.push_back({Expr(func_like.dim(i).min()), Expr(func_like.dim(i).extent())});
     }
 
-    return constant_exterior(Internal::func_like_to_func(func_like), value, object_bounds);
+    return constant_exterior(::Halide::Internal::func_like_to_func(func_like), value, object_bounds);
 }
 template<typename T>
 HALIDE_NO_USER_CODE_INLINE Func constant_exterior(const T &func_like, const Expr &value) {
@@ -127,7 +119,7 @@ HALIDE_NO_USER_CODE_INLINE Func constant_exterior(const T &func_like, const Tupl
                                                   Bounds &&...bounds) {
     Region collected_bounds;
     Internal::collect_region(collected_bounds, std::forward<Bounds>(bounds)...);
-    return constant_exterior(Internal::func_like_to_func(func_like), value, collected_bounds);
+    return constant_exterior(::Halide::Internal::func_like_to_func(func_like), value, collected_bounds);
 }
 template<typename T, typename... Bounds,
          typename std::enable_if<Halide::Internal::all_are_convertible<Expr, Bounds...>::value>::type * = nullptr>
@@ -154,7 +146,7 @@ Func repeat_edge(const Func &source, const Region &bounds);
 
 template<typename T>
 HALIDE_NO_USER_CODE_INLINE Func repeat_edge(const T &func_like, const Region &bounds) {
-    return repeat_edge(Internal::func_like_to_func(func_like), bounds);
+    return repeat_edge(::Halide::Internal::func_like_to_func(func_like), bounds);
 }
 
 template<typename T>
@@ -164,7 +156,7 @@ HALIDE_NO_USER_CODE_INLINE Func repeat_edge(const T &func_like) {
         object_bounds.push_back({Expr(func_like.dim(i).min()), Expr(func_like.dim(i).extent())});
     }
 
-    return repeat_edge(Internal::func_like_to_func(func_like), object_bounds);
+    return repeat_edge(::Halide::Internal::func_like_to_func(func_like), object_bounds);
 }
 // @}
 
@@ -185,7 +177,7 @@ Func repeat_image(const Func &source, const Region &bounds);
 
 template<typename T>
 HALIDE_NO_USER_CODE_INLINE Func repeat_image(const T &func_like, const Region &bounds) {
-    return repeat_image(Internal::func_like_to_func(func_like), bounds);
+    return repeat_image(::Halide::Internal::func_like_to_func(func_like), bounds);
 }
 
 template<typename T>
@@ -195,7 +187,7 @@ HALIDE_NO_USER_CODE_INLINE Func repeat_image(const T &func_like) {
         object_bounds.push_back({Expr(func_like.dim(i).min()), Expr(func_like.dim(i).extent())});
     }
 
-    return repeat_image(Internal::func_like_to_func(func_like), object_bounds);
+    return repeat_image(::Halide::Internal::func_like_to_func(func_like), object_bounds);
 }
 
 /** Impose a boundary condition such that the entire coordinate space is
@@ -216,7 +208,7 @@ Func mirror_image(const Func &source, const Region &bounds);
 
 template<typename T>
 HALIDE_NO_USER_CODE_INLINE Func mirror_image(const T &func_like, const Region &bounds) {
-    return mirror_image(Internal::func_like_to_func(func_like), bounds);
+    return mirror_image(::Halide::Internal::func_like_to_func(func_like), bounds);
 }
 
 template<typename T>
@@ -226,7 +218,7 @@ HALIDE_NO_USER_CODE_INLINE Func mirror_image(const T &func_like) {
         object_bounds.push_back({Expr(func_like.dim(i).min()), Expr(func_like.dim(i).extent())});
     }
 
-    return mirror_image(Internal::func_like_to_func(func_like), object_bounds);
+    return mirror_image(::Halide::Internal::func_like_to_func(func_like), object_bounds);
 }
 
 // @}
@@ -251,7 +243,7 @@ Func mirror_interior(const Func &source, const Region &bounds);
 
 template<typename T>
 HALIDE_NO_USER_CODE_INLINE Func mirror_interior(const T &func_like, const Region &bounds) {
-    return mirror_interior(Internal::func_like_to_func(func_like), bounds);
+    return mirror_interior(::Halide::Internal::func_like_to_func(func_like), bounds);
 }
 
 template<typename T>
@@ -261,7 +253,7 @@ HALIDE_NO_USER_CODE_INLINE Func mirror_interior(const T &func_like) {
         object_bounds.push_back({Expr(func_like.dim(i).min()), Expr(func_like.dim(i).extent())});
     }
 
-    return mirror_interior(Internal::func_like_to_func(func_like), object_bounds);
+    return mirror_interior(::Halide::Internal::func_like_to_func(func_like), object_bounds);
 }
 
 // @}

src/CMakeLists.txt

@@ -67,6 +67,8 @@ set(HEADER_FILES
     FlattenNestedRamps.h
     Float16.h
     Func.h
+    FuncExtras.h
+    FuncTypeChanging.h
     Function.h
     FunctionPtr.h
     FuseGPUThreadLoops.h
@@ -229,6 +231,7 @@ set(SOURCE_FILES
     FlattenNestedRamps.cpp
     Float16.cpp
     Func.cpp
+    FuncTypeChanging.cpp
     Function.cpp
     FuseGPUThreadLoops.cpp
     FuzzFloatStores.cpp

src/FuncExtras.h

@@ -0,0 +1,24 @@
+#ifndef HALIDE_FUNC_EXTRAS_H
+#define HALIDE_FUNC_EXTRAS_H
+
+#include "Func.h"
+#include "Lambda.h"
+
+namespace Halide {
+
+namespace Internal {
+
+inline const Func &func_like_to_func(const Func &func) {
+    return func;
+}
+
+template<typename T>
+inline HALIDE_NO_USER_CODE_INLINE Func func_like_to_func(const T &func_like) {
+    return lambda(_, func_like(_));
+}
+
+}  // namespace Internal
+
+}  // namespace Halide
+
+#endif