[Mosaic TPU] Canonicalize arith.select's condition to vector if other…

… types are vector.

This fixes the failure in elementwise rule of apply vector layout pass.

If the condition scalar is static, it will be simplified to corresponding vector from true value and false value by MLIR.

If the condition scalar is dynamic, we want to use vselect over scf.if anyway. Because latter creates a inner region.

PiperOrigin-RevId: 679764355

jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc

@@ -332,6 +332,25 @@ LogicalResult canonicalize_extract(int hardware_generation, Operation &raw_op) {
   return success();
 }
 
+LogicalResult canonicalize_select(int hardware_generation, Operation &raw_op) {
+  auto op = dyn_cast<arith::SelectOp>(raw_op);
+  if (!isa<VectorType>(op.getType()) ||
+      isa<VectorType>(op.getCondition().getType())) {
+    return success();
+  }
+  // Canonicalize `cond ? vector_a : vector_b` ->
+  // `vector_cond ? vector_a : vector_b`.
+  ImplicitLocOpBuilder builder(op->getLoc(), op.getOperation());
+  auto cond_ty = VectorType::get(cast<VectorType>(op.getType()).getShape(),
+                                 op.getCondition().getType());
+  auto cond = builder.create<vector::BroadcastOp>(cond_ty, op.getCondition());
+  auto new_op = builder.create<arith::SelectOp>(
+      op.getLoc(), cond, op.getTrueValue(), op.getFalseValue());
+  op.replaceAllUsesWith(new_op.getResult());
+  op.erase();
+  return success();
+}
+
 using canonicalize_rule_type =
     std::function<LogicalResult(int hardware_generation, Operation &op)>;
 
@@ -341,7 +360,8 @@ const llvm::StringMap<canonicalize_rule_type> &rules() {
       {vector::ContractionOp::getOperationName(), canonicalize_contraction},
       {vector::ContractionOp::getOperationName(), canonicalize_extract},
       {vector::MultiDimReductionOp::getOperationName(),
-       canonicalize_multi_dim_reduction}};
+       canonicalize_multi_dim_reduction},
+      {arith::SelectOp::getOperationName(), canonicalize_select}};
   return *rules;
 }
 

tests/pallas/tpu_ops_test.py

@@ -208,6 +208,31 @@ def body(x_ref, o_ref):
         ),
     )
 
+  def test_select_with_scalar_condition(self):
+    def kernel(cond, lhs, rhs, out):
+      out[:] = jax.lax.select(cond[0] != 0, lhs[:], rhs[:])
+
+    def run(cond, lhs, rhs):
+      return pl.pallas_call(
+          kernel,
+          out_shape=lhs,
+          grid_spec=pltpu.PrefetchScalarGridSpec(
+              num_scalar_prefetch=0,
+              in_specs=[
+                  pl.BlockSpec(memory_space=pltpu.SMEM),
+                  pl.BlockSpec(memory_space=pltpu.VMEM),
+                  pl.BlockSpec(memory_space=pltpu.VMEM),
+              ],
+          ),
+          name="select_kernel",
+      )(cond, lhs, rhs)
+
+    cond = jnp.array([1], dtype=jnp.int32)
+    lhs = jnp.zeros((8, 128), dtype=jnp.float32)
+    rhs = jnp.ones((8, 128), dtype=jnp.float32)
+
+    assert (run(cond, lhs, rhs) == lhs).all()
+
 
 class OpsInterpretTest(OpsTest):
   INTERPRET = True