Skip to content

Commit 0ffc67d

Browse files
authored
[TOOLS] Fixed bug in AOT compiler (#5771)
The bug was causing the generated code to have CUBIN array that is twice as large as the actual byte count.
1 parent e7457d3 commit 0ffc67d

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

python/triton/tools/compile.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -136,11 +136,12 @@ def constexpr(s):
136136
if hints.get((i, ), None) == 16:
137137
suffix += 'd'
138138
func_name = '_'.join([out_name, sig_hash, suffix])
139-
hex_ = str(binascii.hexlify(ccinfo.asm["cubin"]))[2:-1]
139+
asm = ccinfo.asm["cubin"] # store binary data once
140+
hex_ = str(binascii.hexlify(asm))[2:-1]
140141
params = {
141142
"kernel_name": func_name,
142143
"triton_kernel_name": args.kernel_name,
143-
"bin_size": len(hex_),
144+
"bin_size": len(asm),
144145
"bin_data": ", ".join([f"0x{x}{y}" for x, y in zip(hex_[::2], hex_[1::2])]),
145146
"signature": ", ".join([f"{ty_to_cpp(ty)} {name}" for name, ty in zip(arg_names_not_1, arg_types_not_1)]),
146147
"full_signature": ", ".join([f"{ty_to_cpp(ty)} {name}" for name, ty in zip(arg_names, arg_types)]),

0 commit comments

Comments
 (0)