[Backend] Pass GPUTarget to make_ir (#7689)

This allows us avoid duplicating target parsing logic in different
places, e.g., when in Gluon we don't need to parse it again from backend
options.

Author: antiagainst

python/test/unit/runtime/test_bindings.py

@@ -77,7 +77,7 @@ def walk_fn(op):
     triton._C.libtriton.ir.load_dialects(context)
     backend.load_dialects(context)
 
-    ttir_module = src.make_ir(options, codegen_fns, module_map, context)
+    ttir_module = src.make_ir(target, options, codegen_fns, module_map, context)
     ttir_module.walk(walk_fn)
 
 

python/triton/_filecheck.py

@@ -68,7 +68,7 @@ def run_parser(kernel_fn):
     options = stub_backend.parse_options(options)
     codegen_fns = stub_backend.get_codegen_implementation(options)
     module_map = stub_backend.get_module_map()
-    module = src.make_ir(options, codegen_fns, module_map, context)
+    module = src.make_ir(stub_target, options, codegen_fns, module_map, context)
     assert module.verify()
     return module
 

python/triton/compiler/compiler.py

@@ -77,7 +77,7 @@ def hash(self):
         key = f"{self.fn.cache_key}-{str(self.attrs)}-{sorted_sig}-{constants_key}"
         return hashlib.sha256(key.encode("utf-8")).hexdigest()
 
-    def make_ir(self, options, codegen_fns, module_map, context):
+    def make_ir(self, target: GPUTarget, options, codegen_fns, module_map, context):
         from .code_generator import ast_to_ttir
         return ast_to_ttir(self.fn, self, context=context, options=options, codegen_fns=codegen_fns,
                            module_map=module_map)
@@ -116,7 +116,7 @@ def __init__(self, path, context, backend):
     def hash(self):
         return hashlib.sha256(self.src.encode("utf-8")).hexdigest()
 
-    def make_ir(self, options, codegen_fns, module_map, context):
+    def make_ir(self, target: GPUTarget, options, codegen_fns, module_map, context):
         self.module.context = context
         return self.module
 
@@ -299,7 +299,7 @@ def compile(src, target=None, options=None, _env_vars=None):
     codegen_fns = backend.get_codegen_implementation(options)
     module_map = backend.get_module_map()
     try:
-        module = src.make_ir(options, codegen_fns, module_map, context)
+        module = src.make_ir(target, options, codegen_fns, module_map, context)
     except Exception as e:
         filter_traceback(e)
         raise

python/triton/experimental/gluon/_runtime.py

@@ -1,5 +1,4 @@
 from __future__ import annotations
-import triton
 from triton.compiler.compiler import ASTSource
 from triton.backends.compiler import Language
 from triton.runtime.jit import JITFunction
@@ -16,15 +15,14 @@ def __init__(self, fn, signature, constexprs=None, attrs=None) -> None:
         self.language = Language.GLUON
         self.ext = "ttgir"
 
-    def make_ir(self, options, codegen_fns, module_map, context):
+    def make_ir(self, target, options, codegen_fns, module_map, context):
         from triton.compiler.compiler import make_backend
         from triton.compiler.code_generator import ast_to_ttir
 
         builder = ir.builder(context)
         module = builder.create_module()
 
         # Assign module attributes eagerly, as they are needed to verify layouts
-        target = triton.runtime.driver.active.get_current_target()
         backend = make_backend(target)
         target = backend.get_target_name(options)