Skip to content

Commit

Permalink
WIP
Browse files Browse the repository at this point in the history
  • Loading branch information
LebedevRI committed May 26, 2022
1 parent ad1e7f6 commit 7076ff9
Show file tree
Hide file tree
Showing 11 changed files with 525 additions and 27 deletions.
4 changes: 4 additions & 0 deletions python_bindings/src/PyIROperator.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
34 changes: 13 additions & 21 deletions src/BoundaryConditions.h
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@

#include "Expr.h"
#include "Func.h"
#include "FuncExtras.h"
#include "Lambda.h"

namespace Halide {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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>
Expand All @@ -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) {
Expand All @@ -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>
Expand All @@ -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>
Expand All @@ -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);
}
// @}

Expand All @@ -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>
Expand All @@ -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
Expand All @@ -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>
Expand All @@ -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);
}

// @}
Expand All @@ -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>
Expand All @@ -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);
}

// @}
Expand Down
3 changes: 3 additions & 0 deletions src/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,8 @@ set(HEADER_FILES
FlattenNestedRamps.h
Float16.h
Func.h
FuncExtras.h
FuncTypeChanging.h
Function.h
FunctionPtr.h
FuseGPUThreadLoops.h
Expand Down Expand Up @@ -229,6 +231,7 @@ set(SOURCE_FILES
FlattenNestedRamps.cpp
Float16.cpp
Func.cpp
FuncTypeChanging.cpp
Function.cpp
FuseGPUThreadLoops.cpp
FuzzFloatStores.cpp
Expand Down
24 changes: 24 additions & 0 deletions src/FuncExtras.h
Original file line number Diff line number Diff line change
@@ -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
134 changes: 134 additions & 0 deletions src/FuncTypeChanging.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,134 @@
#include "FuncTypeChanging.h"

namespace Halide {

static bool operator==(const Var &a, const Var &b) {
return a.same_as(b);
}

namespace FuncTypeChanging {

// NOTE: Precondition: `chunk_idx u< num_chunks`!
static Expr get_nth_chunk(const Expr &value, const Type &chunk_type,
Expr chunk_idx, ChunkOrder chunk_order) {
unsigned num_bits_total = value.type().bits();
unsigned num_bits_per_chunk = chunk_type.bits();
unsigned num_chunks = num_bits_total / num_bits_per_chunk;
user_assert(num_bits_total > num_bits_per_chunk &&
num_bits_total % num_bits_per_chunk == 0 && num_chunks > 1)
<< "Input value must evenly partition into several chunks.\n";

std::cerr << "chunk_idx.type() " << chunk_idx.type() << "\n";

Expr low_chunk_idx = chunk_order == ChunkOrder::LowestFirst ?
chunk_idx :
Expr(num_chunks - 1) - chunk_idx;
std::cerr << "low_chunk_idx.type() " << low_chunk_idx.type() << "\n";
Expr num_low_padding_bits = Expr(num_bits_per_chunk) * low_chunk_idx;
Expr chunk_bits = extract_bits(value, num_low_padding_bits,
Expr(num_bits_per_chunk));
return cast(chunk_type, chunk_bits);
}

static Expr concatenate_chunks(std::vector<Expr> chunks,
ChunkOrder chunk_order) {
const Type chunk_type = chunks.front().type();
const unsigned chunk_width = chunk_type.bits();
Type final_type = chunk_type.with_bits(chunk_width * chunks.size()).with_code(halide_type_uint);

if (chunk_order != ChunkOrder::LowestFirst) {
std::reverse(std::begin(chunks), std::end(chunks));
}

Expr res = Internal::make_zero(final_type);
for (size_t chunk_idx = 0; chunk_idx != chunks.size(); ++chunk_idx) {
Expr chunk = chunks[chunk_idx];
Expr wide_chunk = cast(final_type, make_unsigned(chunk)); // zero ext
Expr positioned_chunk = wide_chunk << (chunk_width * chunk_idx);
res = res | positioned_chunk;
}

return res;
}

static Func narrow(const Func &wide_input, const Type &dst_type,
unsigned num_chunks, const Var &dim, const std::string &name,
ChunkOrder chunk_order) {
const std::vector<Var> dims = wide_input.args();
user_assert(count(begin(dims), end(dims), dim) == 1)
<< "Expected dimension " << dim << " to represent "
<< "exactly one function argument!\n";

Expr wide_elt_idx = dim / Expr(num_chunks);
Expr chunk_idx = make_unsigned(dim % Expr(num_chunks));

std::vector<Expr> args;
args.reserve(dims.size());
std::transform(dims.begin(), dims.end(), std::back_inserter(args),
[dim, wide_elt_idx](const Var &input_dim) {
return input_dim.same_as(dim) ? wide_elt_idx : input_dim;
});

Func narrowed(name);
narrowed(dims) = get_nth_chunk(wide_input(args), dst_type,
chunk_idx, chunk_order);

return narrowed;
}

static Func widen(const Func &narrow_input, const Type &dst_type,
int num_chunks, const Var &dim, const std::string &name,
ChunkOrder chunk_order) {
const std::vector<Var> dims = narrow_input.args();
user_assert(count(begin(dims), end(dims), dim) == 1)
<< "Expected dimension " << dim << " to represent "
<< "exactly one function argument!\n";

auto dim_index = std::distance(begin(dims),
std::find(begin(dims), end(dims), dim));

std::vector<Expr> baseline_args;
baseline_args.reserve(dims.size());
std::transform(dims.begin(), dims.end(), std::back_inserter(baseline_args),
[](const Var &input_dim) { return input_dim; });

std::vector<Expr> chunks;
chunks.reserve(num_chunks);
std::generate_n(std::back_inserter(chunks), num_chunks,
[&chunks, baseline_args, dim_index, num_chunks, dim,
narrow_input]() {
int chunk_idx = chunks.size();
std::vector<Expr> args = baseline_args;
args[dim_index] = (num_chunks * dim) + chunk_idx;
return narrow_input(args);
});

Func widened(name);
widened(dims) = concatenate_chunks(chunks, chunk_order);

return widened;
}

Func change_type(const Func &input, const Type &dst_type, const Var &dim,
const std::string &name, ChunkOrder chunk_order) {
const Type &src_type = input.output_type();
int src_width = src_type.bits();
int dst_width = dst_type.bits();
bool is_widening = dst_width > src_width;
auto [min_width, max_width] = std::minmax(src_width, dst_width);
int num_chunks = max_width / min_width;
user_assert(src_type.with_bits(dst_width) == dst_type &&
src_type.is_int_or_uint() && src_width != dst_width &&
max_width % min_width == 0 && num_chunks > 1)
<< "The source type " << src_type << " and destination type "
<< dst_type << " must be similar integer types with different widths, "
<< "larger width must be an integral multiple of the smaller width.\n";

return is_widening ?
widen(input, dst_type, num_chunks, dim, name, chunk_order) :
narrow(input, dst_type, num_chunks, dim, name, chunk_order);
}

} // namespace FuncTypeChanging

} // namespace Halide
54 changes: 54 additions & 0 deletions src/FuncTypeChanging.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
#ifndef HALIDE_FUNC_TYPE_CHANGING_H
#define HALIDE_FUNC_TYPE_CHANGING_H

/** \file
* Support for changing the function's return type by fusing a number of
* consequtive elements, or splitting a single element into parts,
* along a certain dimension.
*/

#include "Func.h"
#include "FuncExtras.h"

namespace Halide {

namespace FuncTypeChanging {

enum class ChunkOrder {
// Example:
// i32 0x0D0C0B0A -> 4xi8 -> { 0x0A, 0x0B, 0x0C, 0x0D }
// i32 0x0D0C0B0A -> 2xi16 -> { 0x0B0A, 0x0D0C }
// 4xi8 { 0x0A, 0x0B, 0x0C, 0x0D } -> i32 -> 0x0D0C0B0A
// 2xi16 { 0x0B0A, 0x0D0C } -> i32 -> 0x0D0C0B0A
// 2xi16 { 0x0D0C, 0x0B0A } -> i32 -> 0x0B0A0D0C
LowestFirst,

// Example:
// i32 0x0D0C0B0A -> 4xi8 -> { 0x0D, 0x0C, 0x0B, 0x0A }
// i32 0x0D0C0B0A -> 2xi16t -> { 0x0D0C, 0x0B0A }
// 4xi8 { 0x0A, 0x0B, 0x0C, 0x0D } -> i32 -> 0x0A0B0C0D
// 2xi16 { 0x0B0A, 0x0D0C } -> i32 -> 0x0B0A0D0C
// 2xi16 { 0x0D0C, 0x0B0A } -> i32 -> 0x0D0C0B0A
HighestFirst,

Default = LowestFirst // DO NOT CHANGE.
};

Func change_type(const Func &input, const Type &dst_type, const Var &dim,
const std::string &name,
ChunkOrder chunk_order = ChunkOrder::Default);

template<typename T>
HALIDE_NO_USER_CODE_INLINE Func
change_type(const T &func_like, const Type &dst_type, const Var &dim,
const std::string &name,
ChunkOrder chunk_order = ChunkOrder::Default) {
return change_type(Internal::func_like_to_func(func_like), dst_type, dim,
name, chunk_order);
}

} // namespace FuncTypeChanging

} // namespace Halide

#endif
Loading

0 comments on commit 7076ff9

Please sign in to comment.