pad non-byte-multiple vectors instead of skipping

lib/Nybbler.cpp

@@ -5,10 +5,13 @@
 //
 // The carrier pattern (spec section 3) is the same for every operation:
 //   1. total = K * N.
-//   2. If total % 8 != 0, skip (leave to the default legalizer; no padding yet).
-//   3. bitcast each operand from <K x iN> to the carrier type <total/8 x i8>.
+//   2. If total % 8 != 0, widen each operand to the next byte boundary by
+//      appending zero lanes (K -> K' with K' * N % 8 == 0).
+//   3. bitcast each (padded) operand from <K' x iN> to the carrier type
+//      <K'*N/8 x i8>.
 //   4. emit the operation's body on the carrier operands.
-//   5. bitcast the result back to <K x iN>, replaceAllUsesWith, erase original.
+//   5. bitcast the result back to <K' x iN>, narrow to the original K lanes
+//      (dropping the pad lanes), replaceAllUsesWith, erase original.
 //
 // Only step 4 differs between operations. So that pattern lives once in the
 // dispatch engine (tryLower), and each operation is a single *handler* function
@@ -61,10 +64,11 @@ bool isNarrowFieldVector(Type *T) {
 namespace {
 
 /// Per-instruction context handed to a CarrierHandler. Carries the live builder
-/// (insertion point already at the original instruction), the original field
-/// vector type \c FieldTy (<K x iN>), the byte carrier type \c CarrierTy
-/// (<total/8 x i8>), the field width \c FieldBits (N), and the original opcode
-/// (so handlers shared across several opcodes, like bitwise, can branch on it).
+/// (insertion point already at the original instruction), the field vector type
+/// \c FieldTy (<K' x iN>, already padded to a byte multiple if the original
+/// lane count wasn't one), the byte carrier type \c CarrierTy (<K'*N/8 x i8>),
+/// the field width \c FieldBits (N), and the original opcode (so handlers
+/// shared across several opcodes, like bitwise, can branch on it).
 struct CarrierOp {
   IRBuilder<> &B;
   FixedVectorType *FieldTy;
@@ -148,24 +152,20 @@ static Value *lowerSub(CarrierOp &Op, ArrayRef<Value *> Ops) {
 
 /// SWAR shl/lshr for i4/i2 via per-field bit-serial conditional shifts.
 ///
-/// Shift-amount model: per-field variable amounts, masked to [0, N-1] via
-/// `& (N-1)`. This deliberately defines results for out-of-range amounts
-/// rather than propagating LLVM IR poison semantics, matching the differential
-/// harness' scalar reference behaviour.
-///
-/// For each power-of-two step s = 1, 2, ..., N/2:
-///   - extract whether each field's amount has bit s set
+/// For each power-of-two step s = 1, 2, ..., 2^(N-1):
+///   - extract whether each field's amount has that bit set
 ///   - conditionally apply a carrier shift by s with a boundary mask
 ///   - blend shifted / unshifted per field
 ///
-/// Shift-amount model: each field's amount is its full N-bit value. Amounts
-/// >= N shift every bit out, yielding 0 -- matching LLVM's scalar reference
-/// (the differential harness' ground truth), which over-shifts to 0 rather
-/// than masking the amount into range.
+/// Shift-amount model: each field's amount is its full N-bit value; amounts
+/// >= N shift every bit out, yielding 0. That out-of-range behaviour is our
+/// choice, not something the scalar reference pins down: `shl/lshr iN x, amt`
+/// with amt >= N is poison, and lli's observed result varies by target and
+/// LLVM build (0 on some, amt % N on others -- it has flipped across CI
+/// updates). The differential kernels therefore mask amounts into [0, N-1]
+/// in-kernel, and only in-range amounts are verified against the reference.
 ///
-/// i1 is special-cased to identity: its only in-range amount is 0, and the
-/// scalar reference (lli interpreting `shl/lshr i1 x, 1`) returns x, so we
-/// match that rather than clearing the field.
+/// i1 is special-cased to identity: its only in-range amount is 0.
 static Value *lowerShift(CarrierOp &Op, ArrayRef<Value *> Ops) {
   IRBuilder<> &B = Op.B;
   unsigned N = Op.FieldBits;
@@ -275,22 +275,54 @@ static bool tryLower(Instruction &I) {
   if (!Handler)
     return false;
 
-  unsigned Total = FieldTy->getNumElements() * FieldBits;
-  if (Total % 8 != 0)
-    return false;
-
   LLVM_DEBUG(dbgs() << "nybbler: lowering " << *BO << "\n");
 
   IRBuilder<> B(BO);
-  auto *CarrierTy = FixedVectorType::get(B.getInt8Ty(), Total / 8);
 
-  Value *LHS = B.CreateBitCast(BO->getOperand(0), CarrierTy);
-  Value *RHS = B.CreateBitCast(BO->getOperand(1), CarrierTy);
+  // Pad non-byte-multiple vectors to the next byte boundary by appending zero
+  // lanes, run the handler on the padded carrier, and drop the pad lanes on
+  // the way back. Zero pad fields are inert for every handler: the field masks
+  // are per-byte splats so they already cover pad fields; add/sub confine
+  // carries/borrows within each field, and a zero field neither generates one
+  // nor lets one escape into the adjacent real field; shift steps operate on
+  // each field independently (the boundary masks confine every carrier shift),
+  // so a pad lane's value can never reach a real lane. Whatever junk the pad
+  // fields do compute is discarded by the narrowing shuffle.
+  unsigned K = FieldTy->getNumElements();
+  unsigned FieldsPerByte = 8 / FieldBits;
+  unsigned PaddedK = (K + FieldsPerByte - 1) / FieldsPerByte * FieldsPerByte;
+
+  Value *LHS = BO->getOperand(0);
+  Value *RHS = BO->getOperand(1);
+  auto *PaddedTy = FieldTy;
+  if (PaddedK != K) {
+    PaddedTy = FixedVectorType::get(FieldTy->getElementType(), PaddedK);
+    // Widen by shuffling with a zero vector: lanes >= K select lane K, i.e.
+    // lane 0 of the zero operand.
+    SmallVector<int, 16> WidenMask;
+    for (unsigned J = 0; J < PaddedK; ++J)
+      WidenMask.push_back(J < K ? J : K);
+    Value *Zero = Constant::getNullValue(FieldTy);
+    LHS = B.CreateShuffleVector(LHS, Zero, WidenMask, "pad.widen.lhs");
+    RHS = B.CreateShuffleVector(RHS, Zero, WidenMask, "pad.widen.rhs");
+  }
+
+  auto *CarrierTy =
+      FixedVectorType::get(B.getInt8Ty(), PaddedK * FieldBits / 8);
 
-  CarrierOp Op{B, FieldTy, CarrierTy, FieldBits, Opcode};
-  Value *Carrier = Handler(Op, {LHS, RHS});
+  Value *CarrierLHS = B.CreateBitCast(LHS, CarrierTy);
+  Value *CarrierRHS = B.CreateBitCast(RHS, CarrierTy);
 
-  Value *Result = B.CreateBitCast(Carrier, FieldTy);
+  CarrierOp Op{B, PaddedTy, CarrierTy, FieldBits, Opcode};
+  Value *Carrier = Handler(Op, {CarrierLHS, CarrierRHS});
+
+  Value *Result = B.CreateBitCast(Carrier, PaddedTy);
+  if (PaddedK != K) {
+    SmallVector<int, 16> NarrowMask;
+    for (unsigned J = 0; J < K; ++J)
+      NarrowMask.push_back(J);
+    Result = B.CreateShuffleVector(Result, NarrowMask, "pad.narrow");
+  }
   BO->replaceAllUsesWith(Result);
   BO->eraseFromParent();
   return true;

test/diff/lshr.ll

@@ -1,19 +1,27 @@
 ; Differential test: per-field lshr on the SWAR carrier must match
 ; LLVM's scalar reference for every input. Driven by tools/diff_runner.py.
+;
+; Shift amounts are masked into [0, N-1] inside each kernel: `lshr iN x, amt`
+; with amt >= N is poison, so the scalar reference's output for it is
+; target/LLVM-build dependent (it has flipped between 0 and amt%N across CI
+; runs). Only in-range amounts have defined semantics to differentiate.
 ; RUN: %python "%diff_runner" --opt "%opt" --lli "%lli" --plugin "%nybbler" "%s" | %FileCheck "%s"
 ; CHECK: ALL PASS
 
 define <64 x i1> @lshr_i1(<64 x i1> %a, <64 x i1> %b) {
-  %r = lshr <64 x i1> %a, %b
+  %amt = and <64 x i1> %b, zeroinitializer
+  %r = lshr <64 x i1> %a, %amt
   ret <64 x i1> %r
 }
 
 define <8 x i2> @lshr_i2(<8 x i2> %a, <8 x i2> %b) {
-  %r = lshr <8 x i2> %a, %b
+  %amt = and <8 x i2> %b, splat (i2 1)
+  %r = lshr <8 x i2> %a, %amt
   ret <8 x i2> %r
 }
 
 define <32 x i4> @lshr_i4(<32 x i4> %a, <32 x i4> %b) {
-  %r = lshr <32 x i4> %a, %b
+  %amt = and <32 x i4> %b, splat (i4 3)
+  %r = lshr <32 x i4> %a, %amt
   ret <32 x i4> %r
 }

test/diff/shl.ll

@@ -1,19 +1,27 @@
 ; Differential test: per-field shl on the SWAR carrier must match
 ; LLVM's scalar reference for every input. Driven by tools/diff_runner.py.
+;
+; Shift amounts are masked into [0, N-1] inside each kernel: `shl iN x, amt`
+; with amt >= N is poison, so the scalar reference's output for it is
+; target/LLVM-build dependent (it has flipped between 0 and amt%N across CI
+; runs). Only in-range amounts have defined semantics to differentiate.
 ; RUN: %python "%diff_runner" --opt "%opt" --lli "%lli" --plugin "%nybbler" "%s" | %FileCheck "%s"
 ; CHECK: ALL PASS
 
 define <64 x i1> @shl_i1(<64 x i1> %a, <64 x i1> %b) {
-  %r = shl <64 x i1> %a, %b
+  %amt = and <64 x i1> %b, zeroinitializer
+  %r = shl <64 x i1> %a, %amt
   ret <64 x i1> %r
 }
 
 define <8 x i2> @shl_i2(<8 x i2> %a, <8 x i2> %b) {
-  %r = shl <8 x i2> %a, %b
+  %amt = and <8 x i2> %b, splat (i2 1)
+  %r = shl <8 x i2> %a, %amt
   ret <8 x i2> %r
 }
 
 define <32 x i4> @shl_i4(<32 x i4> %a, <32 x i4> %b) {
-  %r = shl <32 x i4> %a, %b
+  %amt = and <32 x i4> %b, splat (i4 3)
+  %r = shl <32 x i4> %a, %amt
   ret <32 x i4> %r
 }

test/pad_diff.ll

@@ -0,0 +1,75 @@
+; Differential coverage for the zero-padding path: non-byte-multiple shapes
+; (K*N % 8 != 0) across every lowered op must match LLVM's scalar reference.
+; Representative shapes only -- the full padded matrix belongs to U3-C.
+;
+; Lives at the test root (not diff/) because coverage_check.py treats every
+; diff/*.ll filename as an operation requiring per-width shape tests.
+; RUN: %python "%diff_runner" --opt "%opt" --lli "%lli" --plugin "%nybbler" "%s" | %FileCheck "%s"
+; CHECK: ALL PASS
+
+define <7 x i1> @and_i1_pad(<7 x i1> %a, <7 x i1> %b) {
+  %r = and <7 x i1> %a, %b
+  ret <7 x i1> %r
+}
+
+define <3 x i2> @or_i2_pad(<3 x i2> %a, <3 x i2> %b) {
+  %r = or <3 x i2> %a, %b
+  ret <3 x i2> %r
+}
+
+define <5 x i1> @xor_i1_pad(<5 x i1> %a, <5 x i1> %b) {
+  %r = xor <5 x i1> %a, %b
+  ret <5 x i1> %r
+}
+
+define <13 x i1> @add_i1_pad(<13 x i1> %a, <13 x i1> %b) {
+  %r = add <13 x i1> %a, %b
+  ret <13 x i1> %r
+}
+
+define <5 x i2> @add_i2_pad(<5 x i2> %a, <5 x i2> %b) {
+  %r = add <5 x i2> %a, %b
+  ret <5 x i2> %r
+}
+
+define <3 x i4> @add_i4_pad(<3 x i4> %a, <3 x i4> %b) {
+  %r = add <3 x i4> %a, %b
+  ret <3 x i4> %r
+}
+
+define <7 x i2> @sub_i2_pad(<7 x i2> %a, <7 x i2> %b) {
+  %r = sub <7 x i2> %a, %b
+  ret <7 x i2> %r
+}
+
+define <3 x i4> @sub_i4_pad(<3 x i4> %a, <3 x i4> %b) {
+  %r = sub <3 x i4> %a, %b
+  ret <3 x i4> %r
+}
+
+; Shift amounts are masked into [0, N-1] in-kernel: out-of-range amounts are
+; poison, so the scalar reference's output for them is target/build dependent
+; (see diff/shl.ll).
+define <5 x i2> @shl_i2_pad(<5 x i2> %a, <5 x i2> %b) {
+  %amt = and <5 x i2> %b, splat (i2 1)
+  %r = shl <5 x i2> %a, %amt
+  ret <5 x i2> %r
+}
+
+define <3 x i4> @shl_i4_pad(<3 x i4> %a, <3 x i4> %b) {
+  %amt = and <3 x i4> %b, splat (i4 3)
+  %r = shl <3 x i4> %a, %amt
+  ret <3 x i4> %r
+}
+
+define <7 x i2> @lshr_i2_pad(<7 x i2> %a, <7 x i2> %b) {
+  %amt = and <7 x i2> %b, splat (i2 1)
+  %r = lshr <7 x i2> %a, %amt
+  ret <7 x i2> %r
+}
+
+define <3 x i4> @lshr_i4_pad(<3 x i4> %a, <3 x i4> %b) {
+  %amt = and <3 x i4> %b, splat (i4 3)
+  %r = lshr <3 x i4> %a, %amt
+  ret <3 x i4> %r
+}

test/pad_nonbyte.ll

@@ -0,0 +1,49 @@
+; Vectors whose total bit width is not a multiple of 8 are widened to the next
+; byte boundary with zero lanes, lowered on the padded carrier, and narrowed
+; back to the original lane count. (Replaces skip_nonbyte.ll, which asserted
+; the old behaviour of leaving these to the default legalizer.)
+; RUN: %opt -load-pass-plugin "%nybbler" -passes=nybbler "%s" -S | %FileCheck "%s"
+
+; <3 x i4> = 12 bits -> pad to <4 x i4> = 16 bits, carrier <2 x i8>.
+define <3 x i4> @and_i4_nonbyte(<3 x i4> %a, <3 x i4> %b) {
+; CHECK-LABEL: @and_i4_nonbyte
+; CHECK: shufflevector <3 x i4> %a, <3 x i4> zeroinitializer, <4 x i32>
+; CHECK: shufflevector <3 x i4> %b, <3 x i4> zeroinitializer, <4 x i32>
+; CHECK: bitcast <4 x i4> %{{.*}} to <2 x i8>
+; CHECK: and <2 x i8>
+; CHECK: bitcast <2 x i8> %{{.*}} to <4 x i4>
+; CHECK: shufflevector <4 x i4> %{{.*}}, <3 x i32> <i32 0, i32 1, i32 2>
+; CHECK-NOT: and <3 x i4>
+  %r = and <3 x i4> %a, %b
+  ret <3 x i4> %r
+}
+
+; <5 x i1> = 5 bits -> pad to <8 x i1> = 8 bits, carrier <1 x i8>.
+define <5 x i1> @xor_i1_nonbyte(<5 x i1> %a, <5 x i1> %b) {
+; CHECK-LABEL: @xor_i1_nonbyte
+; CHECK: shufflevector <5 x i1> %a, <5 x i1> zeroinitializer, <8 x i32>
+; CHECK: shufflevector <5 x i1> %b, <5 x i1> zeroinitializer, <8 x i32>
+; CHECK: bitcast <8 x i1> %{{.*}} to <1 x i8>
+; CHECK: xor <1 x i8>
+; CHECK: bitcast <1 x i8> %{{.*}} to <8 x i1>
+; CHECK: shufflevector <8 x i1> %{{.*}}, <5 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4>
+; CHECK-NOT: xor <5 x i1>
+  %r = xor <5 x i1> %a, %b
+  ret <5 x i1> %r
+}
+
+; A field-masked op through the same padding wrapper: the SWAR sub body runs
+; unchanged on the padded <2 x i8> carrier.
+; <5 x i2> = 10 bits -> pad to <8 x i2> = 16 bits, carrier <2 x i8>.
+define <5 x i2> @sub_i2_nonbyte(<5 x i2> %a, <5 x i2> %b) {
+; CHECK-LABEL: @sub_i2_nonbyte
+; CHECK: shufflevector <5 x i2> %a, <5 x i2> zeroinitializer, <8 x i32>
+; CHECK: shufflevector <5 x i2> %b, <5 x i2> zeroinitializer, <8 x i32>
+; CHECK: bitcast <8 x i2> %{{.*}} to <2 x i8>
+; CHECK: sub <2 x i8>
+; CHECK: bitcast <2 x i8> %{{.*}} to <8 x i2>
+; CHECK: shufflevector <8 x i2> %{{.*}}, <5 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4>
+; CHECK-NOT: sub <5 x i2>
+  %r = sub <5 x i2> %a, %b
+  ret <5 x i2> %r
+}