Skip to content

Commit 04d87b5

Browse files
committed
fix hint test on ascend final 3
1 parent 4924d99 commit 04d87b5

File tree

1 file changed

+15
-12
lines changed

1 file changed

+15
-12
lines changed

third_party/ascend/tutorials/hint/test_comment_hint.py

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -26,16 +26,13 @@
2626
This verifies that:
2727
1. #@hint:dot_pad_only_k on tl.load lines generates AnnotationOp with dot_pad_only_k attr in TTIR
2828
2. #@hint:bind_sub_block on for loops generates bind_sub_block attr on scf.for in TTIR
29-
3. The kernel compiles and runs correctly end-to-end
3029
"""
3130

3231
import triton
3332
import triton.language as tl
34-
from triton.compiler.compiler import ASTSource
35-
from triton.backends.ascend.spec.triton.compiler.code_generator import ast_to_ttir
3633
from triton._C.libtriton import ir, ascend
3734
from triton._C.libtriton.ascend import ir as ascend_ir
38-
from triton.backends.ascend.compiler import NPUOptions, min_dot_size
35+
from triton.backends.ascend.compiler import AscendBackend, NPUOptions, min_dot_size
3936
from triton.backends.compiler import GPUTarget
4037

4138

@@ -89,22 +86,28 @@ def matmul_hint_kernel(
8986

9087

9188
# ---------------------------------------------------------------------------
92-
# Helper: compile kernel to TTIR string for IR inspection
89+
# Helper: compile kernel to TTIR string using the full backend pipeline
9390
# ---------------------------------------------------------------------------
9491
def get_ttir_str(kernel_fn, signature, constants):
92+
# Use the ascend backend's compile flow which properly invokes
93+
# the spec ASTSource.make_ir -> ascend ast_to_ttir with hint support
94+
from triton.compiler.compiler import ASTSource
95+
9596
src = ASTSource(kernel_fn, signature, constants)
9697
context = ir.context()
9798
ir.load_dialects(context)
9899
ascend_ir.load_dialects(context)
99100
ascend.load_dialects(context)
101+
100102
options = NPUOptions()
101103
target = GPUTarget("npu", options.arch, 64)
102-
codegen_fns = {"min_dot_size": min_dot_size(target)}
103-
# Apply ascend patch for hint processing
104-
from triton.backends.ascend import _apply_ascend_patch
105-
_apply_ascend_patch()
106-
ttir = ast_to_ttir(kernel_fn, src, context, options, codegen_fns, {})
107-
return str(ttir)
104+
backend = AscendBackend(target)
105+
backend.load_dialects(context)
106+
codegen_fns = backend.get_codegen_implementation()
107+
module_map = backend.get_module_map()
108+
109+
module = src.make_ir(options, codegen_fns, module_map, context)
110+
return str(module)
108111

109112

110113
# ---------------------------------------------------------------------------
@@ -166,4 +169,4 @@ def test_ir_hint_annotations():
166169
# ---------------------------------------------------------------------------
167170
if __name__ == "__main__":
168171
test_ir_hint_annotations()
169-
print("All comment hint tests passed!")
172+
print("All comment hint tests passed!")

0 commit comments

Comments
 (0)