diff --git a/src/CodeGen_Internal.cpp b/src/CodeGen_Internal.cpp index f2e628af3fe5..61ddd8c3cbe9 100644 --- a/src/CodeGen_Internal.cpp +++ b/src/CodeGen_Internal.cpp @@ -795,5 +795,34 @@ void embed_bitcode(llvm::Module *M, const string &halide_command) { } } +Expr lower_concat_bits(const Call *op) { + internal_assert(op->is_intrinsic(Call::concat_bits)); + internal_assert(!op->args.empty()); + + Expr result = make_zero(op->type); + int shift = 0; + for (const Expr &e : op->args) { + result = result | (cast(result.type(), e) << shift); + shift += e.type().bits(); + } + return result; +} + +Expr lower_extract_bits(const Call *op) { + Expr e = op->args[0]; + // Do a shift-and-cast as a uint, which will zero-fill any out-of-range + // bits for us. + if (!e.type().is_uint()) { + e = reinterpret(e.type().with_code(halide_type_uint), e); + } + e = e >> op->args[1]; + e = cast(op->type.with_code(halide_type_uint), e); + if (op->type != e.type()) { + e = reinterpret(op->type, e); + } + e = simplify(e); + return e; +} + } // namespace Internal } // namespace Halide diff --git a/src/CodeGen_Internal.h b/src/CodeGen_Internal.h index faa2a3a4a9d4..8c1a0e1994eb 100644 --- a/src/CodeGen_Internal.h +++ b/src/CodeGen_Internal.h @@ -94,6 +94,12 @@ Expr lower_signed_shift_right(const Expr &a, const Expr &b); /** Reduce a mux intrinsic to a select tree */ Expr lower_mux(const Call *mux); +/** Reduce bit extraction and concatenation to bit ops */ +///@{ +Expr lower_extract_bits(const Call *c); +Expr lower_concat_bits(const Call *c); +///@} + /** Given an llvm::Module, set llvm:TargetOptions information */ void get_target_options(const llvm::Module &module, llvm::TargetOptions &options); diff --git a/src/CodeGen_LLVM.cpp b/src/CodeGen_LLVM.cpp index 30dd7e136657..9e6138b6e4a1 100644 --- a/src/CodeGen_LLVM.cpp +++ b/src/CodeGen_LLVM.cpp @@ -3194,6 +3194,10 @@ void CodeGen_LLVM::visit(const Call *op) { value = codegen(lower_float16_transcendental_to_float32_equivalent(op)); } else if (op->is_intrinsic(Call::mux)) { value = codegen(lower_mux(op)); + } else if (op->is_intrinsic(Call::extract_bits)) { + value = codegen(lower_extract_bits(op)); + } else if (op->is_intrinsic(Call::concat_bits)) { + value = codegen(lower_concat_bits(op)); } else if (op->is_intrinsic()) { Expr lowered = lower_intrinsic(op); if (!lowered.defined()) { diff --git a/src/Deinterleave.cpp b/src/Deinterleave.cpp index e368d851d615..f5840a0074b3 100644 --- a/src/Deinterleave.cpp +++ b/src/Deinterleave.cpp @@ -275,6 +275,16 @@ class Deinterleaver : public IRGraphMutator { return expr; } + Expr give_up_and_shuffle(const Expr &e) { + // Uh-oh, we don't know how to deinterleave this vector expression + // Make llvm do it + std::vector indices; + for (int i = 0; i < new_lanes; i++) { + indices.push_back(starting_lane + lane_stride * i); + } + return Shuffle::make({e}, indices); + } + Expr visit(const Variable *op) override { if (op->type.is_scalar()) { return op; @@ -302,13 +312,7 @@ class Deinterleaver : public IRGraphMutator { lane_stride == 3) { return Variable::make(t, op->name + ".lanes_2_of_3", op->image, op->param, op->reduction_domain); } else { - // Uh-oh, we don't know how to deinterleave this vector expression - // Make llvm do it - std::vector indices; - for (int i = 0; i < new_lanes; i++) { - indices.push_back(starting_lane + lane_stride * i); - } - return Shuffle::make({op}, indices); + return give_up_and_shuffle(op); } } } @@ -325,6 +329,8 @@ class Deinterleaver : public IRGraphMutator { Expr visit(const Reinterpret *op) override { if (op->type.is_scalar()) { return op; + } else if (op->type.bits() != op->value.type().bits()) { + return give_up_and_shuffle(op); } else { Type t = op->type.with_lanes(new_lanes); return Reinterpret::make(t, mutate(op->value)); diff --git a/src/FlattenNestedRamps.cpp b/src/FlattenNestedRamps.cpp index 2b239332feb6..803bd0b85b8f 100644 --- a/src/FlattenNestedRamps.cpp +++ b/src/FlattenNestedRamps.cpp @@ -118,16 +118,54 @@ class FlattenRamps : public IRMutator { } }; +/** Simplify bit concatenation of interleaved loads to vector reinterprets of + * dense loads. Must be done to both vectors and scalars after flattening nested + * ramps, because it can expand a flat ramp into a wider one. */ +class SimplifyConcatBits : public IRMutator { + using IRMutator::visit; + + Expr visit(const Call *op) override { + if (op->is_intrinsic(Call::concat_bits)) { + // Simplify a concat of a load of adjacent bits to a reinterpret of a load of a small vector. + const Load *l0 = op->args[0].as(); + bool ok = true; + const int n = (int)(op->args.size()); + for (int i = 0; ok && i < n; i++) { + const Load *li = op->args[i].as(); + ok &= (li != nullptr); + if (!ok) { + break; + } + const Ramp *r = li->index.as(); + Expr base = r ? r->base : li->index; + ok &= (is_const_one(li->predicate) && + l0->name == li->name && + can_prove(l0->index + i == li->index) && + (r == nullptr || is_const(r->stride, n))); + } + + if (ok) { + internal_assert(l0); + const Ramp *r0 = l0->index.as(); + int new_lanes = (r0 ? r0->lanes : 1) * n; + Expr base = r0 ? r0->base : l0->index; + Expr idx = Ramp::make(base, 1, new_lanes); + return mutate(Reinterpret::make(op->type, Load::make(l0->type.with_lanes(n * l0->type.lanes()), l0->name, idx, l0->image, l0->param, const_true(new_lanes), l0->alignment))); + } + } + + return IRMutator::visit(op); + } +}; + } // namespace Stmt flatten_nested_ramps(const Stmt &s) { - FlattenRamps flatten_ramps; - return flatten_ramps.mutate(s); + return SimplifyConcatBits().mutate(FlattenRamps().mutate(s)); } Expr flatten_nested_ramps(const Expr &e) { - FlattenRamps flatten_ramps; - return flatten_ramps.mutate(e); + return SimplifyConcatBits().mutate(FlattenRamps().mutate(e)); } } // namespace Internal diff --git a/src/IR.cpp b/src/IR.cpp index 740234b8e31f..4b645aea3f87 100644 --- a/src/IR.cpp +++ b/src/IR.cpp @@ -611,12 +611,14 @@ const char *const intrinsic_op_names[] = { "bundle", "call_cached_indirect_function", "cast_mask", + "concat_bits", "count_leading_zeros", "count_trailing_zeros", "debug_to_file", "declare_box_touched", "div_round_to_zero", "dynamic_shuffle", + "extract_bits", "extract_mask_element", "get_user_context", "gpu_thread_barrier", diff --git a/src/IR.h b/src/IR.h index c6085614b59d..edc92f1915c6 100644 --- a/src/IR.h +++ b/src/IR.h @@ -34,7 +34,8 @@ struct Cast : public ExprNode { static const IRNodeType _node_type = IRNodeType::Cast; }; -/** Reinterpret a node as another type, without affecting any of the bits. */ +/** Reinterpret value as another type, without affecting any of the bits + * (on little-endian systems). */ struct Reinterpret : public ExprNode { Expr value; @@ -510,15 +511,26 @@ struct Call : public ExprNode { bitwise_or, bitwise_xor, bool_to_mask, - bundle, // Bundle multiple exprs together temporarily for analysis (e.g. CSE) + + // Bundle multiple exprs together temporarily for analysis (e.g. CSE) + bundle, call_cached_indirect_function, cast_mask, + + // Concatenate bits of the args, with least significant bits as the + // first arg (i.e. little-endian) + concat_bits, count_leading_zeros, count_trailing_zeros, debug_to_file, declare_box_touched, div_round_to_zero, dynamic_shuffle, + + // Extract some contiguous slice of bits from the argument starting at + // the nth bit, counting from the least significant bit, with the number + // of bits determined by the return type. + extract_bits, extract_mask_element, get_user_context, gpu_thread_barrier, @@ -562,7 +574,9 @@ struct Call : public ExprNode { shift_right, signed_integer_overflow, size_of_halide_buffer_t, - sorted_avg, // Compute (arg[0] + arg[1]) / 2, assuming arg[0] < arg[1]. + + // Compute (arg[0] + arg[1]) / 2, assuming arg[0] < arg[1]. + sorted_avg, strict_float, stringify, undef, diff --git a/src/IROperator.cpp b/src/IROperator.cpp index 4693060a8d45..69ea7450b6e0 100644 --- a/src/IROperator.cpp +++ b/src/IROperator.cpp @@ -2654,4 +2654,18 @@ Expr gather(const std::vector &args) { return make_scatter_gather(args); } +Expr extract_bits(Type t, const Expr &e, const Expr &lsb) { + return Internal::Call::make(t, Internal::Call::extract_bits, {e, lsb}, Internal::Call::Intrinsic); +} + +Expr concat_bits(const std::vector &e) { + user_assert(!e.empty()) << "concat_bits requires at least one argument\n"; + user_assert((e.size() & (e.size() - 1)) == 0) << "concat_bits received " << e.size() << " arguments, which is not a power of two.\n"; + Type t = e[0].type(); + for (size_t i = 1; i < e.size(); i++) { + user_assert(e[i].type() == t) << "All arguments to concat_bits must have the same type\n"; + } + return Internal::Call::make(t.with_bits(t.bits() * (int)e.size()), Internal::Call::concat_bits, e, Internal::Call::Intrinsic); +} + } // namespace Halide diff --git a/src/IROperator.h b/src/IROperator.h index ed0b11bb4fef..048998448c75 100644 --- a/src/IROperator.h +++ b/src/IROperator.h @@ -1561,6 +1561,54 @@ Expr gather(const Expr &e, Args &&...args) { } // @} +/** Extract a contiguous subsequence of the bits of 'e', starting at the bit + * index given by 'lsb', where zero is the least-significant bit, returning a + * value of type 't'. Any out-of-range bits requested are filled with zeros. + * + * extract_bits is especially useful when one wants to load a small vector of a + * wide type, and treat it as a larger vector of a smaller type. For example, + * loading a vector of 32 uint8 values from a uint32 Func can be done as + * follows: +\code +f8(x) = extract_bits(f32(x/4), 8*(x%4)); +f8.align_bounds(x, 4).vectorize(x, 32); +\endcode + * Note that the align_bounds call is critical so that the narrow Exprs are + * aligned to the wider Exprs. This makes the x%4 term collapse to a + * constant. If f8 is an output Func, then constraining the min value of x to be + * a known multiple of four would also be sufficient, e.g. via: +\code +f8.output_buffer().dim(0).set_min(0); +\endcode + * + * See test/correctness/extract_concat_bits.cpp for a complete example. */ +// @{ +Expr extract_bits(Type t, const Expr &e, const Expr &lsb); + +template +Expr extract_bits(const Expr &e, const Expr &lsb) { + return extract_bits(type_of(), e, lsb); +} +// @} + +/** Given a number of Exprs of the same type, concatenate their bits producing a + * single Expr of the same type code of the input but with more bits. The + * number of arguments must be a power of two. + * + * concat_bits is especially useful when one wants to treat a Func containing + * values of a narrow type as a Func containing fewer values of a wider + * type. For example, the following code reinterprets vectors of 32 uint8 values + * as a vector of 8 uint32s: + * +\code +f32(x) = concat_bits({f8(4*x), f8(4*x + 1), f8(4*x + 2), f8(4*x + 3)}); +f32.vectorize(x, 8); +\endcode + * + * See test/correctness/extract_concat_bits.cpp for a complete example. + */ +Expr concat_bits(const std::vector &e); + } // namespace Halide #endif diff --git a/src/Simplify_Call.cpp b/src/Simplify_Call.cpp index a1ff4c5130fe..540d7ebc7dff 100644 --- a/src/Simplify_Call.cpp +++ b/src/Simplify_Call.cpp @@ -760,6 +760,8 @@ Expr Simplify::visit(const Call *op, ExprInfo *bounds) { debug(2) << "Simplifier: unhandled PureExtern: " << op->name; } else if (op->is_intrinsic(Call::signed_integer_overflow)) { clear_bounds_info(bounds); + } else if (op->is_intrinsic(Call::concat_bits) && op->args.size() == 1) { + return mutate(op->args[0], bounds); } // No else: we want to fall thru from the PureExtern clause. diff --git a/src/Simplify_Shuffle.cpp b/src/Simplify_Shuffle.cpp index 02f78db260b2..35622aee9c4e 100644 --- a/src/Simplify_Shuffle.cpp +++ b/src/Simplify_Shuffle.cpp @@ -1,4 +1,5 @@ #include "Deinterleave.h" +#include "IROperator.h" #include "Simplify_Internal.h" namespace Halide { @@ -191,6 +192,28 @@ Expr Simplify::visit(const Shuffle *op, ExprInfo *bounds) { } } } + + // Try to collapse an interleave of a series of extract_bits into a vector reinterpret. + if (const Call *extract = new_vectors[0].as()) { + if (extract->is_intrinsic(Call::extract_bits) && + is_const_zero(extract->args[1])) { + int n = (int)new_vectors.size(); + Expr base = extract->args[0]; + bool can_collapse = base.type().bits() == n * op->type.bits(); + for (int i = 1; can_collapse && i < n; i++) { + const Call *c = new_vectors[i].as(); + if (!(c->is_intrinsic(Call::extract_bits) && + is_const(c->args[1], i * op->type.bits()) && + equal(base, c->args[0]))) { + can_collapse = false; + } + } + if (can_collapse) { + return Reinterpret::make(op->type, base); + } + } + } + } else if (op->is_concat()) { // Bypass concat of a single vector (identity shuffle) if (new_vectors.size() == 1) { diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index 62a93e25a982..10a108553231 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -95,6 +95,7 @@ tests(GROUPS correctness extern_reorder_storage.cpp extern_sort.cpp extern_stage_on_device.cpp + extract_concat_bits.cpp failed_unroll.cpp fast_trigonometric.cpp fibonacci.cpp @@ -195,6 +196,7 @@ tests(GROUPS correctness loop_level_generator_param.cpp lossless_cast.cpp lots_of_loop_invariants.cpp + low_bit_depth_noise.cpp make_struct.cpp many_dimensions.cpp many_small_extern_stages.cpp diff --git a/test/correctness/extract_concat_bits.cpp b/test/correctness/extract_concat_bits.cpp new file mode 100644 index 000000000000..20e11a75e1ce --- /dev/null +++ b/test/correctness/extract_concat_bits.cpp @@ -0,0 +1,152 @@ +#include "Halide.h" + +using namespace Halide; + +class CountOps : public Internal::IRMutator { + Expr visit(const Internal::Reinterpret *op) override { + std::cerr << Expr(op) << " " << op->type.lanes() << " " << op->value.type().lanes() << "\n"; + if (op->type.lanes() != op->value.type().lanes()) { + std::cerr << "Got one\n"; + reinterprets++; + } + return Internal::IRMutator::visit(op); + } + + Expr visit(const Internal::Call *op) override { + if (op->is_intrinsic(Internal::Call::concat_bits)) { + concats++; + } else if (op->is_intrinsic(Internal::Call::extract_bits)) { + extracts++; + } + return Internal::IRMutator::visit(op); + } + +public: + int extracts = 0, concats = 0, reinterprets = 0; +}; + +int main(int argc, char **argv) { + for (bool vectorize : {false, true}) { + // Reinterpret an array of a wide type as a larger array of a smaller type + Func f, g; + Var x; + + f(x) = cast(x); + + // Reinterpret to a narrower type. + g(x) = extract_bits(f(x / 4), 8 * (x % 4)); + + f.compute_root(); + + if (vectorize) { + f.vectorize(x, 8); + // The align_bounds directive is critical so that the x%4 term above collapses. + g.align_bounds(x, 4).vectorize(x, 32); + + // An alternative to the align_bounds call: + // g.output_buffer().dim(0).set_min(0); + } + + CountOps counter; + g.add_custom_lowering_pass(&counter, nullptr); + + Buffer out = g.realize({1024}); + std::cerr << counter.extracts << " " << counter.reinterprets << " " << counter.concats << "\n"; + + if (vectorize) { + if (counter.extracts > 0) { + printf("Saw an unwanted extract_bits call in lowered code\n"); + return -1; + } else if (counter.reinterprets == 0) { + printf("Did not see a vector reinterpret in lowered code\n"); + return -1; + } + } + + for (uint32_t i = 0; i < (uint32_t)out.width(); i++) { + uint8_t correct = (i / 4) >> (8 * (i % 4)); + if (out(i) != correct) { + printf("out(%d) = %d instead of %d\n", i, out(i), correct); + return -1; + } + } + } + + for (bool vectorize : {false, true}) { + // Reinterpret an array of a narrow type as a smaller array of a wide type + Func f, g; + Var x; + + f(x) = cast(x); + + g(x) = concat_bits({f(4 * x), f(4 * x + 1), f(4 * x + 2), f(4 * x + 3)}); + + f.compute_root(); + + if (vectorize) { + f.vectorize(x, 32); + g.vectorize(x, 8); + } + + CountOps counter; + g.add_custom_lowering_pass(&counter, nullptr); + + Buffer out = g.realize({64}); + + if (counter.concats > 0) { + printf("Saw an unwanted concat_bits call in lowered code\n"); + return -1; + } else if (counter.reinterprets == 0) { + printf("Did not see a vector reinterpret in lowered code\n"); + return -1; + } + + for (int i = 0; i < 64; i++) { + for (int b = 0; b < 4; b++) { + uint8_t correct = i * 4 + b; + uint8_t result = (out(i) >> (b * 8)) & 0xff; + if (result != correct) { + printf("out(%d) byte %d = %d instead of %d\n", i, b, result, correct); + return -1; + } + } + } + } + + // Also test cases that aren't expected to fold into reinterprets + { + Func f; + Var x("x"); + f(x) = cast(x); + + auto check = [&](const Expr &a, const Expr &b) { + Func g; + g(x) = cast(a == b); + Buffer out = g.realize({1024}); + for (int i = 0; i < out.width(); i++) { + if (out(i) == 0) { + std::cerr << "Mismatch between: " << a << " and " << b << " when x == " << i << "\n"; + exit(-1); + } + } + }; + + // concat_bits is little-endian + check(concat_bits({f(x), cast(37)}), cast(f(x)) + (37 << 16)); + check(concat_bits({cast(0), f(x), cast(0), cast(0)}), cast(UInt(64), f(x)) << 16); + + // extract_bits is equivalent to right shifting and then casting to a narrower type + check(extract_bits(f(x), 3), cast(f(x) >> 3)); + + // Extract bits zero-fills out-of-range bits + check(extract_bits(f(x), 3), f(x) >> 3); + check(extract_bits(f(x), 8), (f(x) >> 8) & 0xff); + check(extract_bits(f(x), -1), cast(f(x)) << 1); + + // MSB of the mantissa of an ieee float + check(extract_bits(cast(f(x)), 15), cast(reinterpret(cast(f(x))) >> 15)); + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/low_bit_depth_noise.cpp b/test/correctness/low_bit_depth_noise.cpp new file mode 100644 index 000000000000..d9d24d60d7b0 --- /dev/null +++ b/test/correctness/low_bit_depth_noise.cpp @@ -0,0 +1,46 @@ +#include "Halide.h" + +using namespace Halide; + +int main(int argc, char **argv) { + // Halide only provides 32-bit noise functions, which are overkill for + // generating low bit-depth noise (e.g. for dithering). This test shows how + // to generate 8-bit noise by slicing out bytes from 32-bit noise. + Var x; + + Func noise; + noise(x) = random_uint(); + + Func noise8; + noise8(x) = extract_bits(noise(x / 4), 8 * (x % 4)); + + Func in16; + in16(x) = cast(x); + + Func dithered; + dithered(x) = cast((in16(x) + noise8(x)) >> 8); + + in16.compute_root(); + dithered.compute_root().vectorize(x, 16, TailStrategy::RoundUp); + noise8.compute_at(dithered, x).vectorize(x); + + // To keep things aligned: + dithered.output_buffer().dim(0).set_min(0); + + Buffer out = dithered.realize({1 << 15}); + + uint32_t sum = 0, correct_sum = 0; + for (int i = 0; i < out.width(); i++) { + sum += out(i); + correct_sum += i; + } + correct_sum = (correct_sum + 128) >> 8; + + if (std::abs((double)sum - correct_sum) / correct_sum > 1e-4) { + printf("Suspiciously large relative difference between the sum of the dithered values and the full-precision sum: %d vs %d\n", sum, correct_sum); + return -1; + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/simplify.cpp b/test/correctness/simplify.cpp index 2771c26c1206..aa10815663d2 100644 --- a/test/correctness/simplify.cpp +++ b/test/correctness/simplify.cpp @@ -2348,6 +2348,10 @@ int main(int argc, char **argv) { Evaluate::make(0)); } + { + check(concat_bits({x}), x); + } + // Check a bounds-related fuzz tester failure found in issue https://github.com/halide/Halide/issues/3764 check(Let::make("b", 105, 336 / max(cast(cast(Variable::make(Int(32), "b"))), 38) + 29), 32);