[inductor] use triu ref instead of lowering (#96040) (#96462)

atalman · Natalia Gimelshein · web-flow · commit c263bd43e8e8 · 2023-03-09T17:42:00.000-05:00
Fixes #95958 Generated code is functionally identical with ref and lowering, only minor differences Pull Request resolved: #96040 Approved by: https://github.com/jansel Co-authored-by: Natalia Gimelshein <ngimel@fb.com>
diff --git a/test/inductor/test_torchinductor_opinfo.py b/test/inductor/test_torchinductor_opinfo.py
@@ -448,6 +448,7 @@ def wrapper_set_seed(op, *args, **kwargs):
     "mT",
     "mH",
     "rsub",
+    "triu",
 }
 
 
diff --git a/torch/_decomp/__init__.py b/torch/_decomp/__init__.py
@@ -308,6 +308,7 @@ def core_aten_decompositions() -> Dict[OpOverload, Callable]:
             aten.trace,
             aten.transpose.int,
             aten.tril.default,
+            aten.triu.default,
             aten.unfold,
             aten.unfold_backward,
             aten.upsample_bilinear2d,
diff --git a/torch/_inductor/lowering.py b/torch/_inductor/lowering.py
@@ -1505,30 +1505,6 @@ def fn(index):
     )
 
 
-@register_lowering(aten.triu)
-def triu(x, diagonal=0):
-    x_loader = x.make_loader()
-    dtype = x.get_dtype()
-
-    def inner_fn(index):
-        *_, i, j = index
-        return ops.where(
-            ops.ge(
-                ops.index_expr(j - i - diagonal, torch.int32),
-                ops.constant(0, torch.int32),
-            ),
-            x_loader(index),
-            ops.constant(0, dtype),
-        )
-
-    return Pointwise.create(
-        device=x.get_device(),
-        dtype=dtype,
-        inner_fn=inner_fn,
-        ranges=list(x.get_size()),
-    )
-
-
 @register_lowering(aten.select_scatter, type_promotion_kind=None)
 def select_scatter(x, src, dim: int, index: int):
     assert x.get_dtype() == src.get_dtype()

Original file line number	Diff line number	Diff line change
`@@ -448,6 +448,7 @@ def wrapper_set_seed(op, args, *kwargs):`
`448`	`448`	`"mT",`
`449`	`449`	`"mH",`
`450`	`450`	`"rsub",`
	`451`	`+ "triu",`
`451`	`452`	`}`
`452`	`453`
`453`	`454`