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

Fixes #95958
Generated code is functionally identical with ref and lowering, only minor differences

Pull Request resolved: https://github.com/pytorch/pytorch/pull/96040
Approved by: https://github.com/jansel

Co-authored-by: Natalia Gimelshein <ngimel@fb.com>

setup.py

@@ -1024,6 +1024,7 @@ def main():
         'typing-extensions',
         'sympy',
         'networkx',
+        'jinja2',
     ]
 
     extras_require = {

test/inductor/test_torchinductor_opinfo.py

@@ -448,6 +448,7 @@ inductor_all_samples = {
     "mT",
     "mH",
     "rsub",
+    "triu",
 }
 
 

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,

torch/_inductor/lowering.py

@@ -1505,30 +1505,6 @@ def iota(
     )
 
 
-@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()

torch/_ops.py

@@ -22,11 +22,14 @@ def dl_open_guard():
     Context manager to set the RTLD_GLOBAL dynamic linker flag while we open a
     shared library to load custom operators.
     """
-    if _SET_GLOBAL_FLAGS:
-        old_flags = sys.getdlopenflags()
-        sys.setdlopenflags(old_flags | ctypes.RTLD_GLOBAL)
-    yield
-    if _SET_GLOBAL_FLAGS:
+    if not _SET_GLOBAL_FLAGS:
+        yield
+        return
+    old_flags = sys.getdlopenflags()
+    sys.setdlopenflags(old_flags | ctypes.RTLD_GLOBAL)
+    try:
+        yield
+    finally:
         sys.setdlopenflags(old_flags)
 
 

torch/utils/tensorboard/summary.py

@@ -380,7 +380,7 @@ def make_histogram(values, bins, max_bins=None):
         limits = new_limits
 
     # Find the first and the last bin defining the support of the histogram:
-    cum_counts = np.cumsum(np.greater(counts, 0, dtype=np.int32))
+    cum_counts = np.cumsum(np.greater(counts, 0))
     start, end = np.searchsorted(cum_counts, [0, cum_counts[-1] - 1], side="right")
     start = int(start)
     end = int(end) + 1