[SPEC] Fix spec language | compiler | triton | gluon init (#790)

* [SPEC] Fix spec language init

* [SPEC] Fix spec compiler & triton init

* [SPEC] Fix spec gluon init

* [SPEC] Fix spec gluon

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,21 +21,19 @@
 )
 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
 
 from . import language
 from . import testing
 from . import tools
 
-# flagtree backend language extension
-from .flagtree_spec import spec
-
-spec("init_language")
-
 must_use_result = language.core.must_use_result
 
+# flagtree backend specialization
+spec("triton_extend_globals", globals())
+
 __all__ = [
     "AsyncCompileMode",
     "autotune",
@@ -49,7 +49,6 @@
     "JITFunction",
     "KernelInterface",
     "language",
-    "max_shared_mem",
     "MockTensor",
     "must_use_result",
     "next_power_of_2",

python/triton/_filecheck.py

@@ -10,15 +10,14 @@
 from triton.experimental.gluon._runtime import GluonASTSource
 from triton.runtime.jit import create_function_from_signature
 from triton._C.libtriton import ir
+from triton.flagtree_spec import spec
 
 # ===-----------------------------------------------------------------------===#
 # filecheck_test
 # ===-----------------------------------------------------------------------===#
 
 # Stub target for testing the frontend.
-# flagtree backend path specialization
-from triton.flagtree_spec import spec
-
+# flagtree backend specialization
 stub_target = spec("spec_get_stub_target")
 if not stub_target:
     stub_target = GPUTarget("cuda", 100, 32)

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/experimental/gluon/__init__.py

@@ -1,16 +1,15 @@
-# flagtree backend path specialization
 from triton.flagtree_spec import spec_path
 
+# flagtree backend path specialization
 spec_path(__path__)
 
-from triton._flagtree_backend import FLAGTREE_BACKEND
+from ._runtime import constexpr_function, jit
+from triton.language.core import must_use_result
 from . import nvidia
-# flagtree backend path specialization
-if FLAGTREE_BACKEND == "hcu":
+from triton._flagtree_backend import FLAGTREE_BACKEND
+if FLAGTREE_BACKEND == "hcu":  # flagtree hcu
     from . import hcu
 else:
     from . import amd
-from ._runtime import constexpr_function, jit
-from triton.language.core import must_use_result
 
 __all__ = ["constexpr_function", "jit", "must_use_result", "nvidia", "hcu" if FLAGTREE_BACKEND == "hcu" else "amd"]

python/triton/experimental/gluon/amd/__init__.py

@@ -1,5 +1,6 @@
 from triton.flagtree_spec import spec_path
 
+# flagtree backend path specialization
 spec_path(__path__)
 
 from . import gfx1250

python/triton/experimental/gluon/language/__init__.py

@@ -1,7 +1,3 @@
-from triton.flagtree_spec import spec_path
-
-spec_path(__path__)
-
 from ._core import (
     base_value,
     base_type,

python/triton/experimental/gluon/language/nvidia/__init__.py

@@ -1,5 +1,6 @@
 from triton.flagtree_spec import spec_path
 
+# flagtree backend path specialization
 spec_path(__path__)
 
 from . import blackwell

python/triton/experimental/gluon/nvidia/__init__.py

@@ -1,5 +1,6 @@
 from triton.flagtree_spec import spec_path
 
+# flagtree backend path specialization
 spec_path(__path__)
 
 from . import hopper

python/triton/language/__init__.py

@@ -1,9 +1,9 @@
 """isort:skip_file"""
 # Import order is significant here.
 
-# 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 . import math
@@ -115,7 +115,6 @@
     store,
     sub,
     tensor,
-    to_tensor,
     trans,
     tuple,
     tuple_type,
@@ -143,6 +142,9 @@
 )
 from . import target_info
 
+# flagtree backend specialization
+spec("language_extend_globals", globals())
+
 __all__ = [
     "PropagateNan",
     "TRITON_MAX_TENSOR_NUMEL",
@@ -270,7 +272,6 @@
     "swizzle2d",
     "target_info",
     "tensor",
-    "to_tensor",
     "topk",
     "trans",
     "tuple",

python/triton/runtime/__init__.py

@@ -1,6 +1,6 @@
-# flagtree backend path specialization
 from triton.flagtree_spec import spec_path
 
+# flagtree backend path specialization
 spec_path(__path__)
 
 from .autotuner import (Autotuner, Config, Heuristics, autotune, heuristics)