Fix example based on triton-lang#2701. (triton-lang#3340)

openxla · Mar 12, 2024 · 098b660 · 098b660
1 parent d04f288
commit 098b660
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/python/examples/copy_strided.py b/python/examples/copy_strided.py
@@ -1,5 +1,6 @@
 import triton
 import triton.language as tl
+import triton.compiler as tc
 
 
 # triton kernel
@@ -14,5 +15,11 @@ def kernel(X, stride_xm,  #
     tl.store(Zs, tl.load(Xs))
 
 
-ret = triton.compile(kernel, signature="*fp32,i32,*fp32,i32", constants={"BLOCK_M": 64, "BLOCK_N": 64})
+src = tc.ASTSource(
+    fn=kernel,
+    constants={"BLOCK_M": 64, "BLOCK_N": 64},
+    signature="*fp32,i32,*fp32,i32",
+)
+
+ret = triton.compile(src)
 print(ret.asm["ttgir"])