From 715d79d9bafbb5630da88e7bb02f0fac9c425d3e Mon Sep 17 00:00:00 2001
From: Philippe Tillet <phil@openai.com>
Date: Fri, 24 Jan 2025 14:06:46 -0800
Subject: [PATCH] [FRONTEND] assert identifier legality (#5697)

---
 python/triton/compiler/code_generator.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/python/triton/compiler/code_generator.py b/python/triton/compiler/code_generator.py
index e5cc0132f5a8..a98ccd71e933 100644
--- a/python/triton/compiler/code_generator.py
+++ b/python/triton/compiler/code_generator.py
@@ -20,6 +20,13 @@
 from .errors import (CompilationError, CompileTimeAssertionFailure, UnsupportedLanguageConstruct)
 
 
+def check_identifier_legality(name, type):
+    pattern = r'^[a-zA-Z_][a-zA-Z0-9_]*$'
+    if not re.match(pattern, name):
+        raise CompilationError(f"invalid {type} identifier: {name}", name)
+    return name
+
+
 def mangle_ty(ty):
     if ty.is_tuple():
         return 'T' + '_'.join(map(mangle_ty, ty.types)) + 'T'
@@ -276,6 +283,9 @@ def __init__(self, context, prototype, gscope, function_name, jit_fn: JITFunctio
 
         self.lscope = {}
         self.jit_fn = jit_fn
+        # TODO: we currently generate illegal names for non-kernel functions involving constexprs!
+        if is_kernel:
+            function_name = check_identifier_legality(function_name, "function")
         self.function_name = function_name
         self.is_kernel = is_kernel
         self.cur_node = None