Fused ZeroWhere operator

[Quafadas]

What's missing is a general masked-zeroing primitive. The backward relu is just: "zero elements of A where corresponding elements of B ≤ 0." That's a masked store — a fundamental SIMD concept that appears everywhere in numerical computing (sparse ops, boundary conditions, thresholding filters, etc.).

The composable API could look like:

// Then relu backward composes naturally:
On SIMD hardware this is a single fused pass — load both arrays, compare, blend with zero, store. No boolean array allocated, one pass over memory. The JVM Vector API has VectorMask.blend() which maps directly to this.

The current composition dz1 := (z1 > 0) is correct and readable, but it forces two passes + a temporary boolean allocation. A zeroWhere! gives users the fused version when they need the performance, while keeping := + > for cases where readability matters more.

Write up a separate (SIMD) version for the JVM, and then loop versions for js and jvm.

A single test suite for all three platforms.

Here's an implementation and testing plan.

zeroWhere! — fused masked zeroing

Motivation

The backward ReLU pattern is:

val mask = z1 > 0          // allocate Array[Boolean], one SIMD pass
dz1 *:*= mask              // second SIMD pass, reads boolean array

This does two passes over memory and allocates a temporary Array[Boolean](numel). At batch=512, hidden=128, that's a 65K boolean allocation per step.

A fused zeroWhere! zeros elements of vec where the corresponding element of other satisfies a comparison against a threshold — one pass, zero allocations.

// zeros elements of vec where other(i) <= 0.0f
vec.`zeroWhere!`(other, 0.0f, ComparisonOp.LE)

// relu backward composes as:
dz1.raw.`zeroWhere!`(z1.raw, 0.0f, ComparisonOp.LE)

API surface

Array-level operations (the core)

Location: floatarrays.scala (JVM), floatarrays.scala (JS), floatarrays.scala (Native).
Repeat for doublearrays.scala / JsNativeDoubleArrays.scala on each platform.

extension (vec: Array[Float])

  /** Zeros elements of `vec` in place where the corresponding element
    * of `other` satisfies the comparison against `threshold`.
    *
    * Requires `vec.length == other.length`.
    *
    * Example: vec.`zeroWhere!`(other, 0.0f, ComparisonOp.LE)
    *   zeros vec(i) wherever other(i) <= 0.0f
    */
  inline def `zeroWhere!`(
      other: Array[Float],
      threshold: Float,
      inline op: ComparisonOp
  ): Unit

  /** Non-mutating variant. Returns a new array. */
  inline def zeroWhere(
      other: Array[Float],
      threshold: Float,
      inline op: ComparisonOp
  ): Array[Float]

Identical signatures for Array[Double].

ComparisonOp

A simple enum to avoid exposing JVM-only VectorOperators.Comparison in cross-platform code.

Location: shared source (vecxt/src/), e.g. a new file ComparisonOp.scala or added to an existing shared definitions file.

enum ComparisonOp:
  case LT, LE, GT, GE, EQ, NE

Matrix-level operations (optional, lower priority)

These delegate to the array versions for contiguous layouts, or fall back to element-wise loops for views. Only needed if callers want to write matrix.zeroWhere!(otherMatrix, ...) directly instead of operating on .raw.

Implementation

JVM (floatarrays.scala / doublearrays.scala)

Use the Vector API. The inner loop is a fused compare-blend-store:

inline def `zeroWhere!`(
    other: Array[Float],
    threshold: Float,
    inline op: ComparisonOp
): Unit =
  assert(vec.length == other.length)
  val zero = FloatVector.zero(spf)
  val thresh = FloatVector.broadcast(spf, threshold)
  var i = 0

  while i < spf.loopBound(vec.length) do
    val values = FloatVector.fromArray(spf, vec, i)
    val cmp = FloatVector.fromArray(spf, other, i)
    // mask is true where condition is met → zero those elements
    val mask = inline op match
      case ComparisonOp.LE => cmp.compare(VectorOperators.LE, thresh)
      case ComparisonOp.LT => cmp.compare(VectorOperators.LT, thresh)
      case ComparisonOp.GE => cmp.compare(VectorOperators.GE, thresh)
      case ComparisonOp.GT => cmp.compare(VectorOperators.GT, thresh)
      case ComparisonOp.EQ => cmp.compare(VectorOperators.EQ, thresh)
      case ComparisonOp.NE => cmp.compare(VectorOperators.NE, thresh)
    // where mask is true → zero; where false → keep original
    values.blend(zero, mask).intoArray(vec, i)
    i += spfl
  end while

  // scalar tail
  while i < vec.length do
    val hit = inline op match
      case ComparisonOp.LE => other(i) <= threshold
      case ComparisonOp.LT => other(i) < threshold
      case ComparisonOp.GE => other(i) >= threshold
      case ComparisonOp.GT => other(i) > threshold
      case ComparisonOp.EQ => other(i) == threshold
      case ComparisonOp.NE => other(i) != threshold
    if hit then vec(i) = 0.0f
    i += 1
  end while
end `zeroWhere!`

Key: values.blend(zero, mask) replaces only the masked lanes with zero. Single load of vec, single load of other, single store. No boolean array.

JS and Native (floatarrays.scala)

Plain while loop — same logic, no SIMD:

inline def `zeroWhere!`(
    other: Array[Float],
    threshold: Float,
    inline op: ComparisonOp
): Unit =
  assert(vec.length == other.length)
  var i = 0
  while i < vec.length do
    val hit = inline op match
      case ComparisonOp.LE => other(i) <= threshold
      case ComparisonOp.LT => other(i) < threshold
      case ComparisonOp.GE => other(i) >= threshold
      case ComparisonOp.GT => other(i) > threshold
      case ComparisonOp.EQ => other(i) == threshold
      case ComparisonOp.NE => other(i) != threshold
    if hit then vec(i) = 0.0f
    i += 1
  end while
end `zeroWhere!`

Non-mutating variant

inline def zeroWhere(...): Array[Float] =
  vec.clone().tap(_.`zeroWhere!`(other, threshold, op))

Files to modify

File	Change
vecxt/src/ComparisonOp.scala (new)	enum ComparisonOp
vecxt/src-jvm/floatarrays.scala	Add zeroWhere! / zeroWhere with SIMD
vecxt/src-jvm/doublearrays.scala	Same for Array[Double]
vecxt/src-js/floatarrays.scala	Add zeroWhere! / zeroWhere (while loop)
vecxt/src-js/doublearrays.scala	Same for Array[Double]
vecxt/src-native/floatarrays.scala	Add zeroWhere! / zeroWhere (while loop)
vecxt/src-native/doublearrays.scala	Same for Array[Double]
vecxt/test/src/	Cross-platform tests (see below)
vecxt/test/src-jvm/	JVM-specific SIMD edge-case tests

Note: if JS and Native share implementations via src-js-native/, use that for DRY. Check whether JsNativeDoubleArrays.scala / equivalent float file exist in src-js-native/ and add there if so.

Testing

All tests in the cross-platform test directory (vecxt/test/src/) unless noted.

1. Basic correctness — Float

test("zeroWhere! zeros elements where other <= threshold (Float)"):
  val vec   = Array[Float](1.0f, 2.0f, 3.0f, 4.0f, 5.0f)
  val other = Array[Float](0.0f, -1.0f, 1.0f, 0.0f, 2.0f)
  vec.`zeroWhere!`(other, 0.0f, ComparisonOp.LE)
  // other(0)=0 LE 0 → true  → zero
  // other(1)=-1 LE 0 → true → zero
  // other(2)=1 LE 0 → false → keep
  // other(3)=0 LE 0 → true  → zero
  // other(4)=2 LE 0 → false → keep
  assertArrayEquals(vec, Array[Float](0.0f, 0.0f, 3.0f, 0.0f, 5.0f))

2. Basic correctness — Double

Same pattern with Array[Double].

3. Non-mutating variant returns new array, source unchanged

test("zeroWhere returns new array without mutating source (Float)"):
  val vec   = Array[Float](1.0f, 2.0f, 3.0f)
  val other = Array[Float](-1.0f, 1.0f, -1.0f)
  val result = vec.zeroWhere(other, 0.0f, ComparisonOp.LE)
  assertArrayEquals(result, Array[Float](0.0f, 2.0f, 0.0f))
  assertArrayEquals(vec, Array[Float](1.0f, 2.0f, 3.0f))  // unchanged

4. All six comparison operators

test("zeroWhere! respects all ComparisonOp variants (Float)"):
  // other = [1, 2, 3], threshold = 2
  // LT: other < 2  → [true, false, false]  → zero index 0
  // LE: other <= 2  → [true, true, false]   → zero index 0,1
  // GT: other > 2   → [false, false, true]  → zero index 2
  // GE: other >= 2  → [false, true, true]   → zero index 1,2
  // EQ: other == 2  → [false, true, false]  → zero index 1
  // NE: other != 2  → [true, false, true]   → zero index 0,2
  val other = Array[Float](1.0f, 2.0f, 3.0f)
  for (op, expected) <- Seq(
    (ComparisonOp.LT, Array(0.0f, 20.0f, 30.0f)),
    (ComparisonOp.LE, Array(0.0f, 0.0f, 30.0f)),
    (ComparisonOp.GT, Array(10.0f, 20.0f, 0.0f)),
    (ComparisonOp.GE, Array(10.0f, 0.0f, 0.0f)),
    (ComparisonOp.EQ, Array(10.0f, 0.0f, 30.0f)),
    (ComparisonOp.NE, Array(0.0f, 20.0f, 0.0f)),
  ) do
    val vec = Array[Float](10.0f, 20.0f, 30.0f)
    vec.`zeroWhere!`(other, 2.0f, op)
    assertArrayEquals(vec, expected, clue = s"failed for $op")

5. SIMD boundary: length not a multiple of vector species

test("zeroWhere! handles non-SIMD-aligned lengths (Float)"):
  // spfl is typically 8 or 16; use length = spfl + 3 to exercise scalar tail
  val n = 19  // likely not a multiple of any SIMD width
  val vec   = Array.tabulate[Float](n)(i => (i + 1).toFloat)
  val other = Array.tabulate[Float](n)(i => if i % 3 == 0 then -1.0f else 1.0f)
  vec.`zeroWhere!`(other, 0.0f, ComparisonOp.LE)
  var i = 0
  while i < n do
    if i % 3 == 0 then assertEquals(vec(i), 0.0f, s"index $i should be zeroed")
    else assertEquals(vec(i), (i + 1).toFloat, s"index $i should be kept")
    i += 1
  end while

6. Empty arrays

test("zeroWhere! on empty arrays is a no-op (Float)"):
  val vec   = Array.empty[Float]
  val other = Array.empty[Float]
  vec.`zeroWhere!`(other, 0.0f, ComparisonOp.LE)  // should not throw
  assertEquals(vec.length, 0)

7. All elements zeroed / none zeroed

test("zeroWhere! zeros all elements when all satisfy condition"):
  val vec   = Array[Float](1.0f, 2.0f, 3.0f)
  val other = Array[Float](-1.0f, -2.0f, -3.0f)
  vec.`zeroWhere!`(other, 0.0f, ComparisonOp.LT)
  assertArrayEquals(vec, Array[Float](0.0f, 0.0f, 0.0f))

test("zeroWhere! keeps all elements when none satisfy condition"):
  val vec   = Array[Float](1.0f, 2.0f, 3.0f)
  val other = Array[Float](10.0f, 20.0f, 30.0f)
  vec.`zeroWhere!`(other, 0.0f, ComparisonOp.LT)
  assertArrayEquals(vec, Array[Float](1.0f, 2.0f, 3.0f))

8. Threshold edge cases

test("zeroWhere! handles NaN in other array (Float)"):
  val vec   = Array[Float](1.0f, 2.0f, 3.0f)
  val other = Array[Float](Float.NaN, 1.0f, Float.NaN)
  // NaN comparisons are always false (IEEE 754)
  vec.`zeroWhere!`(other, 0.0f, ComparisonOp.LE)
  assertArrayEquals(vec, Array[Float](1.0f, 2.0f, 3.0f))  // NaN LE 0 = false

test("zeroWhere! handles threshold at Float boundaries"):
  val vec   = Array[Float](1.0f, 2.0f)
  val other = Array[Float](Float.NegativeInfinity, Float.PositiveInfinity)
  vec.`zeroWhere!`(other, 0.0f, ComparisonOp.LE)
  // -Inf LE 0 → true → zero; +Inf LE 0 → false → keep
  assertArrayEquals(vec, Array[Float](0.0f, 2.0f))

9. Self-referential: vec == other (zeroWhere on self)

test("zeroWhere! works when vec and other are the same array (Float)"):
  val arr = Array[Float](-2.0f, 3.0f, -1.0f, 4.0f)
  arr.`zeroWhere!`(arr, 0.0f, ComparisonOp.LT)
  // clampMin(0) via zeroWhere — negatives zeroed, positives kept
  assertArrayEquals(arr, Array[Float](0.0f, 3.0f, 0.0f, 4.0f))

10. JVM-only: large array (exercises full SIMD loop + tail)

Location: vecxt/test/src-jvm/

test("zeroWhere! on large array matches scalar reference"):
  val n = 100_003
  val rng = new java.util.Random(42)
  val vec   = Array.tabulate[Float](n)(_ => rng.nextFloat() * 10 - 5)
  val other = Array.tabulate[Float](n)(_ => rng.nextFloat() * 10 - 5)
  val expected = vec.clone()
  // scalar reference
  var i = 0
  while i < n do
    if other(i) <= 0.0f then expected(i) = 0.0f
    i += 1
  end while
  vec.`zeroWhere!`(other, 0.0f, ComparisonOp.LE)
  assertArrayEquals(vec, expected)

Performance validation

After implementation, update the MNIST benchmark to use the fused version:

// Before (two-pass):
val dz1Check_ = z1_ > 0
dz1_ *:*= dz1Check_

// After (one-pass, no allocation):
dz1_.raw.`zeroWhere!`(z1_.raw, 0.0f, ComparisonOp.LE)

Expected improvement: eliminates one 65K boolean allocation + one full SIMD pass. The bwd_07_mask_multiply benchmark (which includes the matmul + mask) should show measurable improvement. Consider adding a dedicated micro-benchmark:

@Benchmark
def bwd_07_mask_multiply_fused(bh: Blackhole): Unit =
  val dz = (dz2 @@ w2T)
  dz.raw.`zeroWhere!`(z1.raw, 0.0f, ComparisonOp.LE)  // wait — z1 has different shape than dz

Note: the fused version only works when both arrays share the same contiguous memory layout and length. For the MNIST backward pass, dz1 is (bs, 128) and z1 is (bs, 128) — same shape, both col-major from matmul. So .raw arrays are directly compatible.