[SPEC] Fix spec compiler & triton init

Author: zhzhcookie

python/triton/__init__.py

@@ -19,7 +19,7 @@
 )
 from .runtime.jit import constexpr_function, jit
 from .runtime._async_compile import AsyncCompileMode, FutureKernel
-from .compiler import compile, CompilationError, max_shared_mem
+from .compiler import compile, CompilationError
 from .errors import TritonError
 from .runtime._allocation import set_allocator
 
@@ -29,6 +29,10 @@
 
 must_use_result = language.core.must_use_result
 
+# flagtree backend specialization
+from triton.flagtree_spec import spec
+spec("triton_extend_globals", globals())
+
 __all__ = [
     "AsyncCompileMode",
     "autotune",
@@ -44,7 +48,6 @@
     "JITFunction",
     "KernelInterface",
     "language",
-    "max_shared_mem",
     "MockTensor",
     "must_use_result",
     "next_power_of_2",

python/triton/compiler/__init__.py

@@ -1,12 +1,15 @@
 # flagtree backend path specialization
 from triton.flagtree_spec import spec_path
-
 spec_path(__path__)
 
-from .compiler import CompiledKernel, ASTSource, IRSource, compile, make_backend, LazyDict, get_cache_key, max_shared_mem
+from .compiler import CompiledKernel, ASTSource, IRSource, compile, make_backend, LazyDict, get_cache_key
 from .errors import CompilationError
 
+# flagtree backend specialization
+from triton.flagtree_spec import spec
+spec("compiler_extend_globals", globals())
+
 __all__ = [
     "compile", "make_backend", "ASTSource", "IRSource", "CompiledKernel", "CompilationError", "LazyDict",
-    "get_cache_key", "max_shared_mem"
+    "get_cache_key"
 ]

python/triton/language/__init__.py

@@ -1,14 +1,10 @@
 """isort:skip_file"""
 # Import order is significant here.
 
-from triton.flagtree_spec import spec_path, spec
-
 # flagtree backend path specialization
+from triton.flagtree_spec import spec_path
 spec_path(__path__)
 
-# flagtree backend specialization
-spec("language_extend_globals", globals())
-
 from . import math
 from . import extra
 from ..backends import language_extensions as ext
@@ -145,6 +141,10 @@
 )
 from . import target_info
 
+# flagtree backend specialization
+from triton.flagtree_spec import spec
+spec("language_extend_globals", globals())
+
 __all__ = [
     "PropagateNan",
     "TRITON_MAX_TENSOR_NUMEL",

third_party/mthreads/backend/spec/triton/__init__.py

@@ -1,3 +1,9 @@
-from ._filecheck import spec_get_stub_target
 from ._utils import apply_with_path, _tuple_create
+from ._filecheck import spec_get_stub_target
+from .compiler import compiler_extend_globals
 from .language import language_extend_globals
+
+def triton_extend_globals(globals_dict):
+    # NOTE: Must use absolute path import.
+    from triton.compiler.compiler import max_shared_mem
+    globals_dict["max_shared_mem"] = max_shared_mem

third_party/mthreads/backend/spec/triton/compiler/__init__.py

@@ -0,0 +1,4 @@
+def compiler_extend_globals(globals_dict):
+    # NOTE: Must use absolute path import.
+    from triton.compiler.compiler import max_shared_mem
+    globals_dict["max_shared_mem"] = max_shared_mem