[Pallas TPU] Add WeirdOp to TPU dialect and add lowering for lax.is_finite

PiperOrigin-RevId: 704888940

Author: ayaka14732

jax/_src/pallas/mosaic/lowering.py

@@ -2385,6 +2385,15 @@ def _and_lowering_rule(ctx: LoweringRuleContext, x, y):
 skip_mlir_conversions.add(lax.and_p)
 
 
+def _is_finite_lowering_rule(ctx: LoweringRuleContext, x):
+  out_aval, = ctx.avals_out
+  out_type = aval_to_ir_type(out_aval)
+  return _not_lowering_rule(ctx, tpu.weird(out_type, x))
+
+
+lowering_rules[lax.is_finite_p] = _is_finite_lowering_rule
+
+
 def _or_lowering_rule(ctx: LoweringRuleContext, x, y):
   x, y = _bcast(x, y, *ctx.avals_in, *ctx.avals_out)
   return arith.ori(x, y)

jaxlib/mosaic/dialect/tpu/tpu.td

@@ -473,6 +473,13 @@ def TPU_BitcastVregOp : TPU_Op<"bitcast_vreg", [Pure]> {
   let assemblyFormat = [{ $input attr-dict `:` type($input) `->` type($output) }];
 }
 
+def TPU_WeirdOp : TPU_Op<"weird", [Pure, ElementwiseMappable]> {
+  let arguments = (ins AnyType:$input);  // F32 vector or scalar
+  let results = (outs AnyType:$output);  // I1 vector or scalar
+  let assemblyFormat = [{ $input attr-dict `:` type($input) `->` type($output) }];
+  let hasVerifier = 1;
+}
+
 def TPU_RollVectorsOp : TPU_Op<"roll_vectors", [Pure]> {
   let arguments = (ins Variadic<AnyVectorOfNonZeroRank>:$input);
   let results = (outs AnyVectorOfNonZeroRank:$output);

jaxlib/mosaic/dialect/tpu/tpu_ops.cc

@@ -1087,6 +1087,32 @@ LogicalResult LogOp::verify() {
                       stringifyCoreType(logging_core_type_maybe->value())));
 }
 
+LogicalResult WeirdOp::verify() {
+  const mlir::Type in_type = getInput().getType();
+  if (const auto in_vec_type = dyn_cast<VectorType>(in_type)) {  // Vector case.
+    if (!in_vec_type.getElementType().isF32()) {
+      return emitOpError("Input type must be F32");
+    }
+    const mlir::Type out_type = getResult().getType();
+    const auto out_vec_type = dyn_cast<VectorType>(out_type);
+    if (!out_vec_type) {
+      return emitOpError("Output type must be a vector when input is a vector");
+    }
+    if (!out_vec_type.getElementType().isInteger(1)) {
+      return emitOpError("Output type must be I1");
+    }
+  } else {  // Scalar case.
+    if (!in_type.isF32()) {
+      return emitOpError("Input type must be F32");
+    }
+    const mlir::Type out_type = getResult().getType();
+    if (!out_type.isInteger(1)) {
+      return emitOpError("Output type must be I1 scalar");
+    }
+  }
+  return success();
+}
+
 }  // namespace tpu
 }  // namespace mlir
 

tests/pallas/ops_test.py

@@ -757,6 +757,49 @@ def kernel(x_ref, o_ref):
     expected = lax.erf_inv(x)
     np.testing.assert_array_equal(out, expected)
 
+  IS_FINITE_TEST_VALUES = [
+      -0.2, jnp.inf, -jnp.inf, jnp.nan, 0.0, 1.0, -1.0, 0.5,
+  ]
+
+  def test_is_finite(self):
+    if jtu.test_device_matches(["gpu"]):
+      self.skipTest("Not supported on GPU")
+
+    size = len(self.IS_FINITE_TEST_VALUES)
+
+    @functools.partial(
+        self.pallas_call,
+        out_shape=jax.ShapeDtypeStruct((size,), jnp.bool_),
+    )
+    def kernel(x_ref, o_ref):
+      o_ref[...] = lax.is_finite(x_ref[...])
+
+    x = jnp.array(self.IS_FINITE_TEST_VALUES, dtype=jnp.float32)
+    out = kernel(x)
+    expected = lax.is_finite(x)
+    self.assertArraysEqual(out, expected)
+
+  def test_is_finite_scalar(self):
+    if jtu.test_device_matches(["gpu"]):
+      self.skipTest("Not supported on GPU")
+
+    size = len(self.IS_FINITE_TEST_VALUES)
+
+    @functools.partial(
+        self.pallas_call,
+        in_specs=(pl.BlockSpec(memory_space=smem_on_tpu()),),
+        out_specs=pl.BlockSpec(memory_space=smem_on_tpu()),
+        out_shape=jax.ShapeDtypeStruct((size,), jnp.bool_),
+    )
+    def kernel(x_ref, o_ref):
+      for i in range(8):
+        o_ref[i] = jnp.isfinite(x_ref[i])
+
+    x = jnp.array(self.IS_FINITE_TEST_VALUES, dtype=jnp.float32)
+    out = kernel(x)
+    expected = lax.is_finite(x)
+    self.assertArraysEqual(out, expected)
+
   ELEMENTWISE_OPS = [
       (
           [jnp.abs, jnp.negative],