Repository: pytorch
Issue: triton import inside function doesn't work
Labels: triaged,module: custom-operators,upstream triton

### 🐛 Describe the bug

TODO: need to see whether this is an issue without torch.library custom triton kernel wrapping 

Apply this to https://github.com/pytorch/pytorch/pull/132466

```diff
diff --git a/test/functorch/test_ac.py b/test/functorch/test_ac.py
index 3a2e07fa147..8e9b11846d2 100644
--- a/test/functorch/test_ac.py
+++ b/test/functorch/test_ac.py
@@ -11,14 +11,6 @@ from torch.utils._triton import has_triton
 from torch.utils.flop_counter import FlopCounterMode, register_flop_formula


-if has_triton():
-    # note: if we only import triton in the test, the test fails:
-    # def relu_kernel_(inp_ptr, out_ptr, sz, BLOCK_SIZE: tl.constexpr):
-    # NameError('tl is not defined')
-    import triton
-    import triton.language as tl
-
-
 def compile_with_ac(f, memory_budget):
     return torch.compile(f, backend="aot_eager_decomp_partition")

@@ -196,6 +188,9 @@ class MemoryBudgetTest(TestCase):
     # needs CUDA, but this test file all needs CUDA.
     @unittest.skipIf(not has_triton(), "test needs triton")
     def test_custom_triton_kernel(self):
+        import triton
+        import triton.language as tl
+
         @triton.jit
         def relu_kernel_(inp_ptr, out_ptr, sz, BLOCK_SIZE: tl.constexpr):
             pid = tl.program_id(0)
```

and I get this error:

```
Traceback (most recent call last):
  File "/data/users/dberard/pytorch/torch/testing/_internal/common_utils.py", line 2918, in wrapper
    method(*args, **kwargs)
  File "/data/users/dberard/pytorch/test/functorch/test_ac.py", line 255, in test_custom_triton_kernel
    expected = call()
  File "/data/users/dberard/pytorch/test/functorch/test_ac.py", line 253, in call
    return f(x, ws)
  File "/data/users/dberard/pytorch/test/functorch/test_ac.py", line 242, in f
    x = torch.ops.testac.triton_relu(x)
  File "/data/users/dberard/pytorch/torch/_ops.py", line 1127, in __call__
    return self._op(*args, **(kwargs or {}))
  File "/data/users/dberard/pytorch/torch/utils/_device.py", line 79, in __torch_function__
    return func(*args, **kwargs)
  File "/data/users/dberard/pytorch/torch/_ops.py", line 1127, in __call__
    return self._op(*args, **(kwargs or {}))
  File "/data/users/dberard/pytorch/torch/_library/autograd.py", line 111, in autograd_impl
    result = Generated.apply(*args, Metadata(keyset, keyword_only_args))  # type: ignore[attr-defined]
  File "/data/users/dberard/pytorch/torch/autograd/function.py", line 574, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/data/users/dberard/pytorch/torch/_library/autograd.py", line 50, in forward
    result = op.redispatch(keyset & _C._after_autograd_keyset, *args, **kwargs)
  File "/data/users/dberard/pytorch/torch/_ops.py", line 732, in redispatch
    return self._handle.redispatch_boxed(keyset, *args, **kwargs)
  File "/data/users/dberard/pytorch/torch/_library/custom_ops.py", line 323, in backend_impl
    result = self._backend_fns[device_type](*args, **kwargs)
  File "/data/users/dberard/pytorch/torch/_library/custom_ops.py", line 361, in wrapped_fn
    return fn(*args, **kwargs)
  File "/data/users/dberard/pytorch/test/functorch/test_ac.py", line 209, in triton_relu
    torch._library.capture_triton(relu_kernel_)[grid](x, y, sz, BLOCK_SIZE)
  File "/data/users/dberard/pytorch/torch/_higher_order_ops/triton_kernel_wrap.py", line 1019, in __call__
    return tracing_triton_hopifier_singleton.call_triton_kernel(
  File "/data/users/dberard/pytorch/torch/_higher_order_ops/triton_kernel_wrap.py", line 952, in call_triton_kernel
    return self.call_HOP(variable, grids, combined_args_raw, tx)
  File "/data/users/dberard/pytorch/torch/_higher_order_ops/triton_kernel_wrap.py", line 994, in call_HOP
    return triton_kernel_wrapper_mutation(
  File "/data/users/dberard/pytorch/torch/_ops.py", line 437, in __call__
    return wrapper()
  File "/data/users/dberard/pytorch/torch/_dynamo/eval_frame.py", line 629, in _fn
    return fn(*args, **kwargs)
  File "/data/users/dberard/pytorch/torch/_ops.py", line 433, in wrapper
    return self_.dispatch(
  File "/data/users/dberard/pytorch/torch/_ops.py", line 417, in dispatch
    return kernel(*args, **kwargs)
  File "/data/users/dberard/pytorch/torch/_higher_order_ops/triton_kernel_wrap.py", line 557, in triton_kernel_wrapper_mutation_dense
    kernel[grid_fn](**kwargs, **constant_args)
  File "/data/users/dberard/triton/python/triton/runtime/jit.py", line 326, in <lambda>
    return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
  File "/data/users/dberard/triton/python/triton/runtime/jit.py", line 643, in run
    kernel = self.compile(
  File "/data/users/dberard/triton/python/triton/compiler/compiler.py", line 281, in compile
    module = src.make_ir(options, codegen_fns, context)
  File "/data/users/dberard/triton/python/triton/compiler/compiler.py", line 113, in make_ir
    return ast_to_ttir(self.fn, self, context=context, options=options, codegen_fns=codegen_fns)
triton.compiler.errors.CompilationError: at 1:51:
def relu_kernel_(inp_ptr, out_ptr, sz, BLOCK_SIZE: tl.constexpr):
                                                   ^
NameError('tl is not defined')

----------------------------------------------------------------------
Ran 1 test in 5.730s

FAILED (errors=1)
```

### Versions

https://github.com/pytorch/pytorch/pull/132466, A100

cc @bertmaher @int3 @nmacchioni @chenyang78 @embg @peterbell10