[pallas backend] implement complex numbers (pytorch#167947)

Pull Request resolved: pytorch#167947
Approved by: https://github.com/jansel

Author: oulgen

test/inductor/test_pallas.py

@@ -458,6 +458,112 @@ def fn(x, row_indices):
         expected = fn(x, row_indices)
         self.assertEqual(result, expected)
 
+    def test_complex64_mul(self):
+        """Test complex64 multiplication."""
+
+        def fn(a, b):
+            return a * b
+
+        compiled = self._compile(fn)
+
+        a = torch.randn(16, dtype=torch.complex64, device=self.DEVICE)
+        b = torch.randn(16, dtype=torch.complex64, device=self.DEVICE)
+        result = compiled(a, b)
+        expected = fn(a, b)
+        self.assertEqual(result, expected)
+
+    def test_complex_conj(self):
+        """Test complex conjugate."""
+
+        def fn(x):
+            return torch.conj(x)
+
+        compiled = self._compile(fn)
+
+        x = torch.randn(16, dtype=torch.complex64, device=self.DEVICE)
+        result = compiled(x)
+        expected = fn(x)
+        self.assertEqual(result, expected)
+
+    def test_complex_real(self):
+        """Test extracting real part of complex tensor."""
+
+        def fn(x):
+            return torch.real(x)
+
+        compiled = self._compile(fn)
+
+        x = torch.randn(16, dtype=torch.complex64, device=self.DEVICE)
+        result = compiled(x)
+        expected = fn(x)
+        self.assertEqual(result, expected)
+
+    def test_complex_imag(self):
+        """Test extracting imaginary part of complex tensor."""
+
+        def fn(x):
+            return torch.imag(x)
+
+        compiled = self._compile(fn)
+
+        x = torch.randn(16, dtype=torch.complex64, device=self.DEVICE)
+        result = compiled(x)
+        expected = fn(x)
+        self.assertEqual(result, expected)
+
+    def test_complex_abs(self):
+        """Test complex absolute value (magnitude)."""
+
+        def fn(x):
+            return torch.abs(x)
+
+        compiled = self._compile(fn)
+
+        x = torch.randn(16, dtype=torch.complex64, device=self.DEVICE)
+        result = compiled(x)
+        expected = fn(x)
+        self.assertEqual(result, expected)
+
+    def test_complex128_conj(self):
+        """Test complex128 conjugate operation."""
+
+        def fn(x):
+            return torch.conj(x)
+
+        compiled = self._compile(fn)
+
+        x = torch.randn(16, dtype=torch.complex128, device=self.DEVICE)
+        result = compiled(x)
+        expected = fn(x)
+        self.assertEqual(result, expected)
+
+    def test_complex_mul_scalar(self):
+        """Test complex multiplication with scalar."""
+
+        def fn(x):
+            return x * 2.5
+
+        compiled = self._compile(fn)
+
+        x = torch.randn(16, dtype=torch.complex64, device=self.DEVICE)
+        result = compiled(x)
+        expected = fn(x)
+        self.assertEqual(result, expected)
+
+    def test_complex_conj_mul(self):
+        """Test conjugate followed by multiplication."""
+
+        def fn(x, y):
+            return torch.conj(x) * y
+
+        compiled = self._compile(fn)
+
+        x = torch.randn(16, dtype=torch.complex64, device=self.DEVICE)
+        y = torch.randn(16, dtype=torch.complex64, device=self.DEVICE)
+        result = compiled(x, y)
+        expected = fn(x, y)
+        self.assertEqual(result, expected)
+
 
 @unittest.skipUnless(HAS_PALLAS, "requires jax and pallas")
 class PallasTestsCUDA(PallasTestsMixin, TestCase):

torch/_inductor/codegen/pallas.py

@@ -218,6 +218,32 @@ def constant(val, dtype: torch.dtype) -> str:
             return "True" if val else "False"
         return f"jnp.array({val}, dtype={jax_dtype})"
 
+    @staticmethod
+    def real(x: str) -> str:
+        return f"jnp.real({x})"
+
+    @staticmethod
+    def imag(x: str) -> str:
+        return f"jnp.imag({x})"
+
+    @staticmethod
+    def conj(x: str) -> str:
+        return f"jnp.conj({x})"
+
+    @staticmethod
+    def angle(x: str) -> str:
+        return f"jnp.angle({x})"
+
+    @staticmethod
+    def view_as_real(x: str) -> str:
+        """View complex tensor as real tensor with extra dimension."""
+        return f"jnp.stack([jnp.real({x}), jnp.imag({x})], axis=-1)"
+
+    @staticmethod
+    def view_as_complex(x: str) -> str:
+        """View real tensor as complex tensor."""
+        return f"({x}[..., 0] + 1j * {x}[..., 1])"
+
 
 class PallasKernel(SIMDKernel):
     """
@@ -642,6 +668,7 @@ def codegen_kernel(self, name: Optional[str] = None) -> str:  # type: ignore[ove
             import jax
             import jax.numpy as jnp
             from jax.experimental import pallas as pl
+            from torch._inductor.runtime.runtime_utils import torch_dtype_to_jax_runtime
             """
             + (
                 "\n            from jax.experimental.pallas import triton as pltriton"
@@ -822,16 +849,6 @@ def codegen_kernel(self, name: Optional[str] = None) -> str:  # type: ignore[ove
                     )
 
             code.writeline("# Prepare output metadata from PyTorch tensor")
-            code.writeline("# Map PyTorch dtype to JAX dtype")
-            code.writeline("_torch_dtype_to_jax = {")
-            code.writeline(
-                "    torch.float32: jnp.float32, torch.float64: jnp.float64, torch.float16: jnp.float16,"
-            )
-            code.writeline(
-                "    torch.int32: jnp.int32, torch.int64: jnp.int64, torch.int16: jnp.int16, torch.int8: jnp.int8,"
-            )
-            code.writeline("    torch.uint8: jnp.uint8, torch.bool: jnp.bool_,")
-            code.writeline("}")
             code.writeline(
                 "out_shapes = ("
                 + ", ".join([f"tuple({name}.shape)" for name in output_params])
@@ -840,7 +857,10 @@ def codegen_kernel(self, name: Optional[str] = None) -> str:  # type: ignore[ove
             code.writeline(
                 "out_dtypes = ("
                 + ", ".join(
-                    [f"_torch_dtype_to_jax[{name}.dtype]" for name in output_params]
+                    [
+                        f"torch_dtype_to_jax_runtime({name}.dtype)"
+                        for name in output_params
+                    ]
                 )
                 + ",)"
             )

torch/_inductor/runtime/runtime_utils.py

@@ -189,9 +189,43 @@ def compile_mps_shader(source: str) -> Any:
         raise SyntaxError(f"failed to compile {source} with {err.msg}") from err
 
 
+def torch_dtype_to_jax_runtime(dtype: torch.dtype) -> Any:
+    """
+    Map PyTorch dtype to actual JAX dtype object at runtime.
+
+    This helper is used in generated Pallas kernels at runtime to convert
+    PyTorch dtypes to JAX dtype objects (not string representations).
+
+    Args:
+        dtype: PyTorch dtype to convert
+
+    Returns:
+        JAX dtype object (e.g., jnp.float32 object itself)
+    """
+    import jax.numpy as jnp  # pyrefly: ignore [import-error]
+
+    dtype_map = {
+        torch.float32: jnp.float32,
+        torch.float64: jnp.float64,
+        torch.float16: jnp.float16,
+        torch.bfloat16: jnp.bfloat16,
+        torch.int32: jnp.int32,
+        torch.int64: jnp.int64,
+        torch.int16: jnp.int16,
+        torch.int8: jnp.int8,
+        torch.uint8: jnp.uint8,
+        torch.bool: jnp.bool_,
+        torch.complex64: jnp.complex64,
+        torch.complex128: jnp.complex128,
+    }
+    if dtype not in dtype_map:
+        raise ValueError(f"Unsupported dtype for JAX conversion: {dtype}")
+    return dtype_map[dtype]
+
+
 def torch_dtype_to_jax(dtype: torch.dtype) -> str:
     """
-    Map PyTorch dtype to JAX dtype expression.
+    Map PyTorch dtype to JAX dtype expression string.
 
     This helper is used at compile time in codegen to generate
     JAX dtype expressions for Pallas kernels.
@@ -202,16 +236,8 @@ def torch_dtype_to_jax(dtype: torch.dtype) -> str:
     Returns:
         JAX dtype expression as string (e.g., "jnp.float32")
     """
-    dtype_map = {
-        torch.float32: "jnp.float32",
-        torch.float64: "jnp.float64",
-        torch.float16: "jnp.float16",
-        torch.bfloat16: "jnp.bfloat16",
-        torch.int32: "jnp.int32",
-        torch.int64: "jnp.int64",
-        torch.int16: "jnp.int16",
-        torch.int8: "jnp.int8",
-        torch.uint8: "jnp.uint8",
-        torch.bool: "jnp.bool_",
-    }
-    return dtype_map.get(dtype, f"jnp.{dtype}")
+    jax_dtype = torch_dtype_to_jax_runtime(dtype)
+    dtype_name = jax_dtype.__name__
+    if dtype_name == "bool":
+        dtype_name = "bool_"
+    return f"jnp.{dtype_name}"