Skip to content

Instantly share code, notes, and snippets.

@knwng

knwng/triton_hack.py

Last active August 19, 2025 07:36
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save knwng/abc22eb8a73070dfb8bd32029bb57d7a to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save knwng/abc22eb8a73070dfb8bd32029bb57d7a to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. knwng revised this gist Aug 19, 2025. 1 changed file with 3 additions and 3 deletions.
    6 changes: 3 additions & 3 deletions triton_hack.py
    View file
    Original file line number Diff line number Diff line change
    @@ -1,4 +1,4 @@
    # At the end of make_ttgir
    # At the end of make_ttgir, e.g., https://github.com/triton-lang/triton/blob/6ca2dda9bdd331d007d6fab342db5a85f9b23c7d/third_party/amd/backend/compiler.py#L258
    import os
    if "AMD_INSERT_TTGIR" in os.environ.keys():
    fn = os.environ['AMD_INSERT_TTGIR']
    @@ -15,7 +15,7 @@
    mod.context = ctx


    # At the end of make_llir
    # At the end of make_llir, e.g., https://github.com/triton-lang/triton/blob/6ca2dda9bdd331d007d6fab342db5a85f9b23c7d/third_party/amd/backend/compiler.py#L406
    import os
    if "AMD_INSERT_LLVM_IR" in os.environ.keys():
    if ':' in os.environ["AMD_INSERT_LLVM_IR"]:
    @@ -33,7 +33,7 @@
    return file.read()


    # At the end of make_amdgcn
    # At the end of make_amdgcn, e.g. https://github.com/triton-lang/triton/blob/6ca2dda9bdd331d007d6fab342db5a85f9b23c7d/third_party/amd/backend/compiler.py#L428
    import os
    if "AMD_INSERT_AMDGCN" in os.environ.keys():
    if ':' in os.environ["AMD_INSERT_AMDGCN"]:
  2. knwng revised this gist Jul 29, 2025. 1 changed file with 4 additions and 1 deletion.
    5 changes: 4 additions & 1 deletion triton_hack.py
    View file
    Original file line number Diff line number Diff line change
    @@ -1,3 +1,4 @@
    # At the end of make_ttgir
    import os
    if "AMD_INSERT_TTGIR" in os.environ.keys():
    fn = os.environ['AMD_INSERT_TTGIR']
    @@ -14,6 +15,7 @@
    mod.context = ctx


    # At the end of make_llir
    import os
    if "AMD_INSERT_LLVM_IR" in os.environ.keys():
    if ':' in os.environ["AMD_INSERT_LLVM_IR"]:
    @@ -30,7 +32,8 @@
    with open(insert_module_path, "r") as file:
    return file.read()



    # At the end of make_amdgcn
    import os
    if "AMD_INSERT_AMDGCN" in os.environ.keys():
    if ':' in os.environ["AMD_INSERT_AMDGCN"]:
  3. knwng revised this gist Jul 10, 2025. 1 changed file with 17 additions and 0 deletions.
    17 changes: 17 additions & 0 deletions triton_hack.py
    View file
    Original file line number Diff line number Diff line change
    @@ -29,3 +29,20 @@
    raise RuntimeError(f'cannot find llvm ir file to insert. Given: `{insert_module_path}`')
    with open(insert_module_path, "r") as file:
    return file.read()


    import os
    if "AMD_INSERT_AMDGCN" in os.environ.keys():
    if ':' in os.environ["AMD_INSERT_AMDGCN"]:
    kernel_name, insert_module_path = os.environ["AMD_INSERT_AMDGCN"].split(':')
    if kernel_name == metadata["name"]:
    print(f"Replace kernel {kernel_name}'s amdgcn with {insert_module_path}")
    else:
    return amdgcn
    else:
    insert_module_path = os.environ["AMD_INSERT_AMDGCN"]
    print(f"Replace kernel's amdgcn with {insert_module_path}")
    if not os.path.exists(insert_module_path):
    raise RuntimeError(f'cannot find amdgcn file to insert. Given: `{insert_module_path}`')
    with open(insert_module_path, "r") as file:
    return file.read()
  4. knwng revised this gist Jun 29, 2025. 1 changed file with 2 additions and 1 deletion.
    3 changes: 2 additions & 1 deletion triton_hack.py
    View file
    Original file line number Diff line number Diff line change
    @@ -1,3 +1,4 @@
    import os
    if "AMD_INSERT_TTGIR" in os.environ.keys():
    fn = os.environ['AMD_INSERT_TTGIR']
    if ':' in fn:
    @@ -13,7 +14,7 @@
    mod.context = ctx


    import os
    if "AMD_INSERT_LLVM_IR" in os.environ.keys():
    if ':' in os.environ["AMD_INSERT_LLVM_IR"]:
    kernel_name, insert_module_path = os.environ["AMD_INSERT_LLVM_IR"].split(':')
  5. knwng revised this gist May 27, 2025. 1 changed file with 1 addition and 1 deletion.
    2 changes: 1 addition & 1 deletion triton_hack.py
    View file
    Original file line number Diff line number Diff line change
    @@ -23,8 +23,8 @@
    return str(llvm_mod)
    else:
    insert_module_path = os.environ["AMD_INSERT_LLVM_IR"]
    print(f"Replace kernel's llir with {insert_module_path}")
    if not os.path.exists(insert_module_path):
    raise RuntimeError(f'cannot find llvm ir file to insert. Given: `{insert_module_path}`')
    print(f"Replace kernel's llir with {insert_module_path}")
    with open(insert_module_path, "r") as file:
    return file.read()
  6. knwng created this gist Apr 18, 2025.
    30 changes: 30 additions & 0 deletions triton_hack.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,30 @@
    if "AMD_INSERT_TTGIR" in os.environ.keys():
    fn = os.environ['AMD_INSERT_TTGIR']
    if ':' in fn:
    kernel_name, insert_module_path = fn.split(':')
    print(f"Replace kernel {kernel_name}'s ttgir with {insert_module_path}")
    if not mod.has_function(kernel_name):
    return mod
    else:
    insert_module_path = fn
    print(f"Replace kernel's ttgir with {insert_module_path}")
    ctx = mod.context
    mod = ir.parse_mlir_module(insert_module_path, ctx)
    mod.context = ctx



    if "AMD_INSERT_LLVM_IR" in os.environ.keys():
    if ':' in os.environ["AMD_INSERT_LLVM_IR"]:
    kernel_name, insert_module_path = os.environ["AMD_INSERT_LLVM_IR"].split(':')
    if kernel_name == fns[0].name:
    print(f"Replace kernel {kernel_name}'s llir with {insert_module_path}")
    else:
    return str(llvm_mod)
    else:
    insert_module_path = os.environ["AMD_INSERT_LLVM_IR"]
    if not os.path.exists(insert_module_path):
    raise RuntimeError(f'cannot find llvm ir file to insert. Given: `{insert_module_path}`')
    print(f"Replace kernel's llir with {insert_module_path}")
    with open(insert_module_path, "r") as file:
    return file.read()