Integrate Triton up to [b2de88f8](https://github.com/openai/triton/co…

…mmits/b2de88f89c1ff8082c92165535e48dece55da392) PiperOrigin-RevId: 667508787
jax-ml · Aug 28, 2024 · 7ad41f2 · 7ad41f2
1 parent bce82cb
commit 7ad41f2
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/jax_triton/triton_lib.py b/jax_triton/triton_lib.py
@@ -319,6 +319,7 @@ def get_or_create_triton_kernel(
     _triton.ir.load_dialects(context)
     cuda_backend.load_dialects(context)
     codegen_fns = cuda_backend.get_codegen_implementation()
+    module_map = cuda_backend.get_module_map()
 
     module = code_gen.ast_to_ttir(
         fn,
@@ -331,6 +332,7 @@ def get_or_create_triton_kernel(
         options=cuda_options,
         codegen_fns=codegen_fns,
         context=context,
+        module_map=module_map
     )
     ttir = str(module)