fix unittests process replay [pr] (#10947)

This commit is contained in:
qazal
2025-06-24 10:30:23 +03:00
committed by GitHub
parent 7d560dbd75
commit 7a5e4e0bf1
+6 -3
View File
@@ -42,9 +42,12 @@ def replay_kernelize(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str, tuple[
return to_str(new_sink), to_str(ret[big_sink]), (big_sink,)
def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer) -> tuple[str, str, tuple[Any, ...]]:
k2 = Kernel(ast, opts=renderer)
k2.apply_opts(p.applied_opts)
optimized_ast = k2.get_optimized_ast(name_override=p.name)
# only use Kernel class if captured ast isn't already optimized
if ast.arg is None:
k2 = Kernel(ast, opts=renderer)
k2.apply_opts(p.applied_opts)
optimized_ast = k2.get_optimized_ast(name_override=p.name)
else: optimized_ast = ast
p2 = get_program(optimized_ast, renderer)
def to_str(ret:ProgramSpec) -> str: return ret.src
return to_str(p2), to_str(p), (p.ast, renderer, p.applied_opts)