WaveCache.module_op already returns the MLIR module str

--------- Signed-off-by: tyb0807 <[email protected]>
iree-org · Feb 13, 2025 · 97c3f6b · 97c3f6b
1 parent 4ef04d7
commit 97c3f6b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/tests/kernel/wave/attention/evoformer_test.py b/tests/kernel/wave/attention/evoformer_test.py
@@ -140,7 +140,7 @@ def testEvoformerAttentionForward(
         if dump_generated_mlir:
             filename = f"wave_evoformer_{'x'.join(map(str, shape))}.mlir"
             with open(filename, "w") as f:
-                f.write(mb.module_op.get_asm())
+                f.write(mb.module_op)
 
         eps = 1e-2 if output.dtype == torch.float16 else 5e-2
         assert (

diff --git a/tests/kernel/wave/wave_gemm_test.py b/tests/kernel/wave/wave_gemm_test.py
@@ -205,7 +205,7 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
         if test_dump_generated_mlir:
             filename = f"wave_gemm_{'x'.join(map(str, shape))}.mlir"
             with open(filename, "w") as f:
-                f.write(mb.module_op.get_asm())
+                f.write(mb.module_op)
 
         if run_bench:
             if dump_perf is not None:
@@ -352,7 +352,7 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
         if test_dump_generated_mlir:
             filename = f"wave_gemm_{'x'.join(map(str, shape))}.mlir"
             with open(filename, "w") as f:
-                f.write(mb.module_op.get_asm())
+                f.write(mb.module_op)
 
         if run_bench:
             if dump_perf is not None:
@@ -501,7 +501,7 @@ def repeat(acc: tkl.Register[M, N, tkl.i32]) -> tkl.Register[M, N, tkl.i32]:
         if test_dump_generated_mlir:
             filename = f"wave_gemm_{'x'.join(map(str, shape))}.mlir"
             with open(filename, "w") as f:
-                f.write(mb.module_op.get_asm())
+                f.write(mb.module_op)
 
         if run_bench:
             if dump_perf is not None:
@@ -618,7 +618,7 @@ def repeat(acc: tkl.Register[M, N, tkl.i32]) -> tkl.Register[M, N, tkl.i32]:
         if test_dump_generated_mlir:
             filename = f"wave_gemm_{'x'.join(map(str, shape))}_f8.mlir"
             with open(filename, "w") as f:
-                f.write(mb.module_op.get_asm())
+                f.write(mb.module_op)
 
         if run_bench:
             if dump_perf is not None:
@@ -733,7 +733,7 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
         if test_dump_generated_mlir:
             filename = f"wave_gemm_{'x'.join(map(str, shape))}_f8.mlir"
             with open(filename, "w") as f:
-                f.write(mb.module_op.get_asm())
+                f.write(mb.module_op)
 
         if run_bench:
             if dump_perf is not None:
@@ -841,7 +841,7 @@ def repeat(
         if test_dump_generated_mlir:
             filename = f"wave_batched_gemm_{'x'.join(map(str, shape))}.mlir"
             with open(filename, "w") as f:
-                f.write(mb.module_op.get_asm())
+                f.write(mb.module_op)
 
         if run_bench:
             if dump_perf is not None: