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

…mmits/b2de88f89c1ff8082c92165535e48dece55da392) PiperOrigin-RevId: 667508787
jax-ml · Aug 27, 2024 · d947e6b · d947e6b
1 parent bce82cb
commit d947e6b
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)