Generalize vectorized addition to notice that uniform + contiguous is contiguous, not varying.

This is relevant for preserving vector loads when an index is offset.

Author: axch

src/lib/Vectorize.hs

@@ -487,9 +487,13 @@ vectorizePrimOp op = case op of
   BinOp opk arg1 arg2 -> do
     sx@(VVal vx x) <- vectorizeAtom arg1
     sy@(VVal vy y) <- vectorizeAtom arg2
-    let v = case (vx, vy) of (Uniform, Uniform) -> Uniform; _ -> Varying
-    x' <- if vx /= v then ensureVarying sx else return x
-    y' <- if vy /= v then ensureVarying sy else return y
+    let v = case (opk, vx, vy) of
+              (_, Uniform, Uniform) -> Uniform
+              (IAdd, Uniform, Contiguous) -> Contiguous
+              (IAdd, Contiguous, Uniform) -> Contiguous
+              _ -> Varying
+    x' <- if v == Varying then ensureVarying sx else return x
+    y' <- if v == Varying then ensureVarying sy else return y
     VVal v <$> emitOp (BinOp opk x' y')
   MiscOp (CastOp tyArg arg) -> do
     ty <- vectorizeType tyArg

tests/opt-tests.dx

@@ -126,13 +126,13 @@ _ = for i:(Fin 20) j:(Fin 4). ordinal j
 "vectorizing int binary op"
 -- CHECK-LABEL: vectorizing int binary op
 %passes vect
-_ = for i:(Fin 256). (n_to_i32 (ordinal i)) + 1
+_ = for i:(Fin 256). (n_to_i32 (ordinal i)) * 2
 -- CHECK: seq (RawFin 0x10)
 -- CHECK: [[i0:v#[0-9]+]]:<16xInt32> = vbroadcast
 -- CHECK: [[i1:v#[0-9]+]]:<16xInt32> = viota
 -- CHECK: [[i2:v#[0-9]+]]:<16xInt32> = %iadd [[i0]] [[i1]]
--- CHECK: [[ones:v#[0-9]+]]:<16xInt32> = vbroadcast 1
--- CHECK: %iadd [[i2]] [[ones]]
+-- CHECK: [[twos:v#[0-9]+]]:<16xInt32> = vbroadcast 2
+-- CHECK: %imul [[i2]] [[twos]]
 
 "vectorizing float binary op"
 -- CHECK-LABEL: vectorizing float binary op