Fixing aten.view aliasing issue for complex views.

htyu · htyu · commit b284ce905bf1 · 2023-10-23T17:33:52.000-07:00
diff --git a/test/inductor/test_torchinductor.py b/test/inductor/test_torchinductor.py
@@ -3056,7 +3056,7 @@ def fn(x):
             fn,
             (torch.randn([1, 2, 4, 8]).to(dtype=torch.complex64),),
         )
-        self.assertEqual(torch._inductor.metrics.generated_kernel_count, 1)
+        self.assertEqual(torch._inductor.metrics.generated_kernel_count, 0)
 
         class ToComplex(nn.Module):
             def forward(self, x):
diff --git a/torch/_inductor/decomposition.py b/torch/_inductor/decomposition.py
@@ -270,59 +270,16 @@ def angle(x):
 
 
 @register_decomposition([aten.add])
-def add(x, y, **kwargs):
+def add(x, y, *, alpha=None):
     x_is_complex_tensor = torch.is_tensor(x) and x.is_complex()
     y_is_complex_tensor = torch.is_tensor(y) and y.is_complex()
-    if not x_is_complex_tensor and not y_is_complex_tensor:
+    if not x_is_complex_tensor or not y_is_complex_tensor:
         return NotImplemented
-    r = y.real
-    i = y.imag if y_is_complex_tensor else 0.0
-    alpha = kwargs.get("alpha")
+    z = y
     if alpha is not None:
-        r *= alpha
-        i *= alpha
-    r = x.real + r
-    if x_is_complex_tensor:
-        i = x.imag + i
-    complex_type = x.dtype if x_is_complex_tensor else y.dtype
-    return (
-        torch.where(
-            torch.arange(2, device=x.device, dtype=torch.uint8) == 0,
-            r.unsqueeze(-1),
-            i.unsqueeze(-1),
-        )
-        .view(complex_type)
-        .squeeze(-1)
-    )
-
-
-@register_decomposition([aten.real])
-def real(self):
-    if not torch.is_tensor(self) or not self.is_complex():
-        return NotImplemented
-    assert self.is_complex(), "real should only be called on a complex tensor"
-    if self.dtype == torch.complex32:
-        return self.view(torch.float16)[..., ::2]
-    elif self.dtype == torch.complex64:
-        return self.view(torch.float32)[..., ::2]
-    elif self.dtype == torch.complex128:
-        return self.view(torch.float64)[..., ::2]
-    else:
-        raise AssertionError("unsupported complex type")
-
-
-@register_decomposition([aten.imag])
-def imag(self):
-    if not torch.is_tensor(self) or not self.is_complex():
-        return NotImplemented
-    if self.dtype == torch.complex32:
-        return self.view(torch.float16)[..., 1::2]
-    elif self.dtype == torch.complex64:
-        return self.view(torch.float32)[..., 1::2]
-    elif self.dtype == torch.complex128:
-        return self.view(torch.float64)[..., 1::2]
-    else:
-        raise AssertionError("unsupported complex type")
+        z *= y
+    complex_type = torch.promote_types(x.dtype, y.dtype)
+    return (x.view(x.real.dtype) + z.view(y.real.dtype)).view(complex_type)
 
 
 @register_decomposition([aten.conj_physical])
diff --git a/torch/_inductor/ir.py b/torch/_inductor/ir.py
@@ -132,7 +132,7 @@ def _check_tensorbox(nodes):
                     sympy.Symbol,
                     sympy.logic.boolalg.Boolean,
                     Expr,
-                    torch._inductor.ir.ExpandView,
+                    torch._inductor.ir.MultiOutput
                 ),
             ), f"Found {type(nodes)}, which is not a supported top level IR node. See [Note: Inductor IR]"
 
@@ -4390,6 +4390,19 @@ def apply_constraint(self):
         return super().apply_constraint()
 
 
+
+@dataclasses.dataclass
+class ComplexView(FallbackKernel):
+    """View a complex number as two dtyped numbers or vice versa"""
+
+    def should_allocate(self):
+        return False
+
+    def get_alias_names(self):
+        # Signal to codegen that our output buffer isn't safe to reuse
+        return [self.inputs[0].get_name()]
+
+
 @dataclasses.dataclass
 class MultiOutputLayout(IRNode):
     device: torch.device
diff --git a/torch/_inductor/lowering.py b/torch/_inductor/lowering.py
@@ -543,6 +543,8 @@ def _to_dtype_bitcast(x):
 
 @register_lowering(aten.view.dtype, type_promotion_kind=None)
 def _view_dtype(x: TensorBox, dtype: torch.dtype):
+    if dtype.is_complex or x.get_dtype().is_complex:
+        return ir.ComplexView.create(torch.ops.aten.view.dtype, x, dtype)
     return to_dtype_bitcast(x, dtype, copy=True)
 
 
@@ -1602,17 +1604,20 @@ def _warn_complex_not_supported():
 
 # There are some types (CPU) which we accept as input but not as
 # output.
-def unsupported_input_tensor(t: torch._subclasses.FakeTensor):
+def unsupported_input_tensor(t: torch._subclasses.FakeTensor, parent):
     "Do not support reading or writing to this tensor"
     if t.is_complex():
+        # Complex views are supported with IR ComplexView
+        if parent and parent.target == torch.ops.aten.view.dtype:
+            return False
         _warn_complex_not_supported()
         return True
     return False
 
 
-def unsupported_output_tensor(t: torch._subclasses.FakeTensor):
+def unsupported_output_tensor(t: torch._subclasses.FakeTensor, parent=None):
     "Do not support writing tensor but can read from it"
-    if unsupported_input_tensor(t):
+    if unsupported_input_tensor(t, parent):
         return True
     return t.is_cpu and config.disable_cpp_codegen
 
@@ -1626,7 +1631,7 @@ def fallback_node_due_to_unsupported_type(node: torch.fx.Node, allow_cpu_inputs=
     if node.target is aten.lift_fresh_copy.default:
         return False
 
-    def check_skip_condition(node, is_output):
+    def check_skip_condition(node, parent):
         if not isinstance(node, torch.fx.Node):
             return False
 
@@ -1637,21 +1642,21 @@ def check_skip_condition(node, is_output):
             if not isinstance(meta, torch._subclasses.FakeTensor):
                 continue
 
-            if is_output:
-                if unsupported_output_tensor(meta):
+            if parent:
+                if unsupported_output_tensor(meta, parent):
                     return True
             else:
-                if unsupported_input_tensor(meta):
+                if unsupported_input_tensor(meta, None):
                     return True
 
         return False
 
     # only skip codegen if there is a cpu output, not input
     for arg in tree_flatten((node.args, node.kwargs))[0]:
-        if check_skip_condition(arg, is_output=False):
+        if check_skip_condition(arg, node):
             return True
 
-    return check_skip_condition(node, is_output=True)
+    return check_skip_condition(node, None)
 
 
 def make_fallback(op, layout_constraint=None, warn=True):
