[SPEC] Fix spec compiler & triton init

Author: zhzhcookie

python/triton/__init__.py

@@ -1,6 +1,8 @@
 """isort:skip_file"""
 __version__ = '3.6.0'
 
+from triton.flagtree_spec import spec
+
 # ---------------------------------------
 # Note: import order is significant here.
 
@@ -19,7 +21,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 +31,9 @@
 
 must_use_result = language.core.must_use_result
 
+# flagtree backend specialization
+spec("triton_extend_globals", globals())
+
 __all__ = [
     "AsyncCompileMode",
     "autotune",
@@ -44,7 +49,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
+from triton.flagtree_spec import spec_path, spec
 
+# flagtree backend path specialization
 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
+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

@@ -6,9 +6,6 @@
 # flagtree backend path specialization
 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 +142,9 @@
 )
 from . import target_info
 
+# flagtree backend specialization
+spec("language_extend_globals", globals())
+
 __all__ = [
     "PropagateNan",
     "TRITON_MAX_TENSOR_NUMEL",

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

@@ -1,3 +1,10 @@
-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

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

@@ -7,4 +7,4 @@ def language_extend_globals(globals_dict):
     globals_dict["unsqueeze"] = unsqueeze
     globals_dict["_experimental_descriptor_load"] = _experimental_descriptor_load
     globals_dict["_experimental_descriptor_store"] = _experimental_descriptor_store
-    globals_dict["to_tensor"] = to_tensor
+    globals_dict["to_tensor"] = to_tensor