Narrowing/widening the Func along a dimension (#6756)

python_bindings/src/halide/CMakeLists.txt

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

python_bindings/src/halide/halide_/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/halide/halide_/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/halide/halide_/PyHalide.cpp

@@ -11,6 +11,7 @@
 #include "PyExpr.h"
 #include "PyExternFuncArgument.h"
 #include "PyFunc.h"
+#include "PyFuncTypeChanging.h"
 #include "PyIROperator.h"
 #include "PyImageParam.h"
 #include "PyInlineReductions.h"
@@ -49,6 +50,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/halide/halide_/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

@@ -70,6 +70,8 @@ set(HEADER_FILES
  FlattenNestedRamps.h
  Float16.h
  Func.h
+ FuncExtras.h
+ FuncTypeChanging.h
  Function.h
  FunctionPtr.h
  FuseGPUThreadLoops.h
@@ -234,6 +236,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