[FRONTEND] [BC Breaking] Require global variables to be insantiated as constexpr ob… (#5961)

Author: zahimoud

python/test/unit/language/test_compile_errors.py

@@ -300,7 +300,11 @@ def kernel():
         a = CONSTEXPR_ANNOTATED_GLOBAL  # noqa
 
     # No error.
-    triton.compile(triton.compiler.ASTSource(fn=kernel, signature={}, constexprs={}))
+    try:
+        triton.compile(triton.compiler.ASTSource(fn=kernel, signature={}, constexprs={}))
+        assert False, "Using a constexpr annotated global variable should not be allowed"
+    except CompilationError as e:
+        assert "Cannot access global variable" in str(e)
 
 
 CONSTEXPR_GLOBAL = tl.constexpr(42)

python/test/unit/language/test_random.py

@@ -110,7 +110,7 @@ def random_raw(self):
 # Unit Tests
 #####################################
 
-BLOCK: tl.constexpr = 1024
+BLOCK = tl.constexpr(1024)
 
 # test generation of random uint32
 
@@ -144,7 +144,7 @@ def const_kernel(X, N, seed: tl.constexpr):
     # triton result
     x = torch.empty(size, dtype=torch_dtype, device=device)
     N = x.numel()
-    grid = (triton.cdiv(N, BLOCK), )
+    grid = (triton.cdiv(N, BLOCK.value), )
     if const_seed:
         const_kernel[grid](x, N, seed=seed)
     else:
@@ -184,7 +184,7 @@ def const_kernel(X, N, seed: tl.constexpr, dtype: tl.constexpr):
     # triton result
     x = torch.empty(size, dtype=torch.float32, device=device)
     N = x.numel()
-    grid = (triton.cdiv(N, BLOCK), )
+    grid = (triton.cdiv(N, BLOCK.value), )
     if const_seed:
         const_kernel[grid](x, N, seed=seed, dtype=getattr(tl, dtype))
     else:
@@ -238,7 +238,7 @@ def const_kernel(X, N, seed: tl.constexpr, dtype: tl.constexpr):
     # triton result
     x = torch.empty(size, dtype=torch.float32, device=device)
     N = x.numel()
-    grid = (triton.cdiv(N, BLOCK), )
+    grid = (triton.cdiv(N, BLOCK.value), )
     if const_seed:
         const_kernel[grid](x, N, seed=seed, dtype=getattr(tl, dtype))
     else:

python/test/unit/runtime/test_cache.py

@@ -223,7 +223,7 @@ def kernel(X, i: tl.constexpr = GLOBAL_DEFAULT_ARG):
     assert len(kernel.device_caches[device][0]) == 1
 
 
-GLOBAL_VAR: tl.constexpr = 1
+GLOBAL_VAR = tl.constexpr(1)
 
 
 def test_kernel_global_var_change(device):
@@ -263,7 +263,7 @@ def kernel():
     kernel[(1, )]()
 
 
-CONSTEXPR_GLOBAL: tl.constexpr = 42
+CONSTEXPR_GLOBAL = tl.constexpr(42)
 
 
 def test_local_does_not_shadow_global():
@@ -274,9 +274,9 @@ def kernel():
         a = CONSTEXPR_GLOBAL  # noqa
         _, CONSTEXPR_GLOBAL = 0, 0  # noqa
 
-    CONSTEXPR_GLOBAL = 42
+    CONSTEXPR_GLOBAL = tl.constexpr(42)
     kernel[(1, )]()
-    CONSTEXPR_GLOBAL = 43
+    CONSTEXPR_GLOBAL = tl.constexpr(43)
 
     # Error because the `CONSTEXPR_GLOBAL` we're modifying is the same
     # `CONSTEXPR_GLOBAL` that's read inside `kernel`.  (Alternatively, we could
@@ -288,7 +288,7 @@ def kernel():
         kernel[(1, )]()
 
 
-CONFLICTING_GLOBAL: tl.constexpr = 0
+CONFLICTING_GLOBAL = tl.constexpr(0)
 
 
 @triton.jit

python/triton/compiler/code_generator.py

@@ -12,7 +12,7 @@
 from .._C.libtriton import ir
 from ..language import constexpr, semantic, str_to_ty, tensor
 from ..language.core import _unwrap_if_constexpr, nv_tma_desc_type, base_value, base_type
-from ..runtime.jit import _normalize_ty, get_jit_fn_file_line
+from ..runtime.jit import get_jit_fn_file_line
 # ideally we wouldn't need any runtime component
 from ..runtime import JITFunction
 from .._utils import find_paths_if, get_iterable_path, set_iterable_path
@@ -347,9 +347,6 @@ def _is_constexpr_global(self, name):
         if _is_constexpr(val):
             return True
 
-        if a := self.gscope.get("__annotations__", {}).get(name):
-            return _normalize_ty(a) == "constexpr"
-
         return False
 
     def _is_namedtuple(self, val):
@@ -386,8 +383,8 @@ def global_lookup(name: str, absent):
                 textwrap.dedent(f"""\
                 Cannot access global variable {name} from within @jit'ed
                 function. Triton kernels can only access global variables that
-                are annotated as constexpr (`x: triton.language.constexpr = 42`
-                or `x = triton.language.constexpr(42)`).  Alternatively, set the
+                are instanstiated as constexpr (`x = triton.language.constexpr(42)`). Note that this is different from
+                annotating a variable as constexpr (`x: triton.language.constexpr = 42`), which is not supported.  Alternatively, set the
                 envvar TRITON_ALLOW_NON_CONSTEXPR_GLOBALS=1, but we do not
                 promise to support this forever.""").replace("\n", " "))