Add minimal useful implementation of extracting and concatenating bits

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

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);
 

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()) {

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<int> 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<int> 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));

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<Load>();
+            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<Load>();
+                ok &= (li != nullptr);
+                if (!ok) {
+                    break;
+                }
+                const Ramp *r = li->index.as<Ramp>();
+                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<Ramp>();
+                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

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",

src/IR.h

@@ -34,7 +34,8 @@ struct Cast : public ExprNode<Cast> {
     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<Reinterpret> {
     Expr value;
 
@@ -510,15 +511,26 @@ struct Call : public ExprNode<Call> {
         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<Call> {
         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,

src/IROperator.cpp

@@ -2654,4 +2654,18 @@ Expr gather(const std::vector<Expr> &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<Expr> &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

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<uint8_t>(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<typename T>
+Expr extract_bits(const Expr &e, const Expr &lsb) {
+    return extract_bits(type_of<T>(), 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<Expr> &e);
+
 }  // namespace Halide
 
 #endif

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.

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<Call>()) {
+            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<Call>();
+                    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) {

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