You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "junrushao (via GitHub)" <gi...@apache.org> on 2023/06/22 04:26:04 UTC
[GitHub] [tvm] junrushao opened a new pull request, #15141: Scaffolding DLight
junrushao opened a new pull request, #15141:
URL: https://github.com/apache/tvm/pull/15141
(no comment)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #15141: [WIP] Scaffolding DLight
Posted by "junrushao (via GitHub)" <gi...@apache.org>.
junrushao commented on code in PR #15141:
URL: https://github.com/apache/tvm/pull/15141#discussion_r1240091398
##########
tests/python/dlight/test_modules.py:
##########
@@ -0,0 +1,2621 @@
+from tvm.script import ir as I
+from tvm.script import tir as T
+
+# fmt: off
+
+@I.ir_module
+class Decode:
+
+ @T.prim_func
+ def func1(A: T.Buffer((T.int64(320), T.int64(2560)), "uint32"), B: T.Buffer((T.int64(80), T.int64(2560)), "float16"), T_transpose: T.Buffer((T.int64(2560), T.int64(2560)), "float16")):
+ T.func_attr({"op_pattern": 2, "tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ decode = T.alloc_buffer((T.int64(2560), T.int64(2560)), "float16")
+ for i, j in T.grid(T.int64(2560), T.int64(2560)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(A[v_i // T.int64(8), v_j], B[v_i // T.int64(32), v_j])
+ T.writes(decode[v_i, v_j])
+ decode[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(A[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * B[v_i // T.int64(32), v_j]
+ for ax0, ax1 in T.grid(T.int64(2560), T.int64(2560)):
+ with T.block("T_transpose"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(decode[v_ax1, v_ax0])
+ T.writes(T_transpose[v_ax0, v_ax1])
+ T_transpose[v_ax0, v_ax1] = decode[v_ax1, v_ax0]
+
+ @T.prim_func
+ def func2(A: T.Buffer((T.int64(320), T.int64(10240)), "uint32"), B: T.Buffer((T.int64(80), T.int64(10240)), "float16"), T_transpose: T.Buffer((T.int64(10240), T.int64(2560)), "float16")):
+ T.func_attr({"op_pattern": 2, "tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ decode = T.alloc_buffer((T.int64(2560), T.int64(10240)), "float16")
+ for i, j in T.grid(T.int64(2560), T.int64(10240)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(A[v_i // T.int64(8), v_j], B[v_i // T.int64(32), v_j])
+ T.writes(decode[v_i, v_j])
+ decode[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(A[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * B[v_i // T.int64(32), v_j]
+ for ax0, ax1 in T.grid(T.int64(10240), T.int64(2560)):
+ with T.block("T_transpose"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(decode[v_ax1, v_ax0])
+ T.writes(T_transpose[v_ax0, v_ax1])
+ T_transpose[v_ax0, v_ax1] = decode[v_ax1, v_ax0]
+
+ @T.prim_func
+ def func3(A: T.Buffer((T.int64(1280), T.int64(2560)), "uint32"), B: T.Buffer((T.int64(320), T.int64(2560)), "float16"), T_transpose: T.Buffer((T.int64(2560), T.int64(10240)), "float16")):
+ T.func_attr({"op_pattern": 2, "tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ decode = T.alloc_buffer((T.int64(10240), T.int64(2560)), "float16")
+ for i, j in T.grid(T.int64(10240), T.int64(2560)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(A[v_i // T.int64(8), v_j], B[v_i // T.int64(32), v_j])
+ T.writes(decode[v_i, v_j])
+ decode[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(A[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * B[v_i // T.int64(32), v_j]
+ for ax0, ax1 in T.grid(T.int64(2560), T.int64(10240)):
+ with T.block("T_transpose"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(decode[v_ax1, v_ax0])
+ T.writes(T_transpose[v_ax0, v_ax1])
+ T_transpose[v_ax0, v_ax1] = decode[v_ax1, v_ax0]
+
+ @T.prim_func
+ def func4(A: T.Buffer((T.int64(512), T.int64(4096)), "uint32"), B: T.Buffer((T.int64(128), T.int64(4096)), "float16"), T_transpose: T.Buffer((T.int64(4096), T.int64(4096)), "float16")):
+ T.func_attr({"op_pattern": 2, "tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ decode = T.alloc_buffer((T.int64(4096), T.int64(4096)), "float16")
+ for i, j in T.grid(T.int64(4096), T.int64(4096)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(A[v_i // T.int64(8), v_j], B[v_i // T.int64(32), v_j])
+ T.writes(decode[v_i, v_j])
+ decode[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(A[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * B[v_i // T.int64(32), v_j]
+ for ax0, ax1 in T.grid(T.int64(4096), T.int64(4096)):
+ with T.block("T_transpose"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(decode[v_ax1, v_ax0])
+ T.writes(T_transpose[v_ax0, v_ax1])
+ T_transpose[v_ax0, v_ax1] = decode[v_ax1, v_ax0]
+
+ @T.prim_func
+ def func5(A: T.Buffer((T.int64(512), T.int64(11008)), "uint32"), B: T.Buffer((T.int64(128), T.int64(11008)), "float16"), T_transpose: T.Buffer((T.int64(11008), T.int64(4096)), "float16")):
+ T.func_attr({"op_pattern": 2, "tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ decode = T.alloc_buffer((T.int64(4096), T.int64(11008)), "float16")
+ for i, j in T.grid(T.int64(4096), T.int64(11008)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(A[v_i // T.int64(8), v_j], B[v_i // T.int64(32), v_j])
+ T.writes(decode[v_i, v_j])
+ decode[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(A[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * B[v_i // T.int64(32), v_j]
+ for ax0, ax1 in T.grid(T.int64(11008), T.int64(4096)):
+ with T.block("T_transpose"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(decode[v_ax1, v_ax0])
+ T.writes(T_transpose[v_ax0, v_ax1])
+ T_transpose[v_ax0, v_ax1] = decode[v_ax1, v_ax0]
+
+ @T.prim_func
+ def func6(A: T.Buffer((T.int64(1376), T.int64(4096)), "uint32"), B: T.Buffer((T.int64(344), T.int64(4096)), "float16"), T_transpose: T.Buffer((T.int64(4096), T.int64(11008)), "float16")):
+ T.func_attr({"op_pattern": 2, "tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ decode = T.alloc_buffer((T.int64(11008), T.int64(4096)), "float16")
+ for i, j in T.grid(T.int64(11008), T.int64(4096)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(A[v_i // T.int64(8), v_j], B[v_i // T.int64(32), v_j])
+ T.writes(decode[v_i, v_j])
+ decode[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(A[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * B[v_i // T.int64(32), v_j]
+ for ax0, ax1 in T.grid(T.int64(4096), T.int64(11008)):
+ with T.block("T_transpose"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(decode[v_ax1, v_ax0])
+ T.writes(T_transpose[v_ax0, v_ax1])
+ T_transpose[v_ax0, v_ax1] = decode[v_ax1, v_ax0]
+
+
+@I.ir_module
+class DecodeGemv:
+
+ @T.prim_func
+ def func1(lv2515: T.Buffer((T.int64(320), T.int64(50432)), "uint32"), lv2516: T.Buffer((T.int64(80), T.int64(50432)), "float32"), lv705: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float32"), var_matmul_intermediate: T.Buffer((T.int64(1), T.int64(1), T.int64(50432)), "float32")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ var_decode_intermediate = T.alloc_buffer((T.int64(2560), T.int64(50432)))
+ for i, j in T.grid(T.int64(2560), T.int64(50432)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(lv2515[v_i // T.int64(8), v_j], lv2516[v_i // T.int64(32), v_j])
+ T.writes(var_decode_intermediate[v_i, v_j])
+ var_decode_intermediate[v_i, v_j] = T.Cast("float32", T.Cast("float16", T.bitwise_and(T.shift_right(lv2515[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv2516[v_i // T.int64(32), v_j]
+ for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), T.int64(50432), T.int64(2560)):
+ with T.block("matmul"):
+ v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
+ T.reads(lv705[v_i0, v_i1, v_k], var_decode_intermediate[v_k, v_i2])
+ T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
+ with T.init():
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float32(0)
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + lv705[v_i0, v_i1, v_k] * var_decode_intermediate[v_k, v_i2]
+
+ @T.prim_func
+ def func2(lv1363: T.Buffer((T.int64(320), T.int64(2560)), "uint32"), lv1364: T.Buffer((T.int64(80), T.int64(2560)), "float16"), lv2067: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16"), linear_bias192: T.Buffer((T.int64(2560),), "float16"), p_output0_intermediate: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ var_decode_intermediate = T.alloc_buffer((T.int64(2560), T.int64(2560)), "float16")
+ var_matmul_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")
+ for i, j in T.grid(T.int64(2560), T.int64(2560)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(lv1363[v_i // T.int64(8), v_j], lv1364[v_i // T.int64(32), v_j])
+ T.writes(var_decode_intermediate[v_i, v_j])
+ var_decode_intermediate[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv1363[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv1364[v_i // T.int64(32), v_j]
+ for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), T.int64(2560), T.int64(2560)):
+ with T.block("matmul"):
+ v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
+ T.reads(lv2067[v_i0, v_i1, v_k], var_decode_intermediate[v_k, v_i2])
+ T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
+ with T.init():
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0)
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + lv2067[v_i0, v_i1, v_k] * var_decode_intermediate[v_k, v_i2]
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("T_add"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_matmul_intermediate[v_ax0, v_ax1, v_ax2], linear_bias192[v_ax2])
+ T.writes(p_output0_intermediate[v_ax0, v_ax1, v_ax2])
+ p_output0_intermediate[v_ax0, v_ax1, v_ax2] = var_matmul_intermediate[v_ax0, v_ax1, v_ax2] + linear_bias192[v_ax2]
+
+ @T.prim_func
+ def func3(lv1381: T.Buffer((T.int64(320), T.int64(2560)), "uint32"), lv1382: T.Buffer((T.int64(80), T.int64(2560)), "float16"), lv328: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16"), linear_bias195: T.Buffer((T.int64(2560),), "float16"), lv2062: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16"), p_output0_intermediate: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ var_decode_intermediate = T.alloc_buffer((T.int64(2560), T.int64(2560)), "float16")
+ var_matmul_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")
+ var_T_add_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")
+ for i, j in T.grid(T.int64(2560), T.int64(2560)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(lv1381[v_i // T.int64(8), v_j], lv1382[v_i // T.int64(32), v_j])
+ T.writes(var_decode_intermediate[v_i, v_j])
+ var_decode_intermediate[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv1381[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv1382[v_i // T.int64(32), v_j]
+ for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), T.int64(2560), T.int64(2560)):
+ with T.block("matmul"):
+ v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
+ T.reads(lv328[v_i0, v_i1, v_k], var_decode_intermediate[v_k, v_i2])
+ T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
+ with T.init():
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0)
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + lv328[v_i0, v_i1, v_k] * var_decode_intermediate[v_k, v_i2]
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("T_add"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_matmul_intermediate[v_ax0, v_ax1, v_ax2], linear_bias195[v_ax2])
+ T.writes(var_T_add_intermediate[v_ax0, v_ax1, v_ax2])
+ var_T_add_intermediate[v_ax0, v_ax1, v_ax2] = var_matmul_intermediate[v_ax0, v_ax1, v_ax2] + linear_bias195[v_ax2]
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("T_add_1"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_T_add_intermediate[v_ax0, v_ax1, v_ax2], lv2062[v_ax0, v_ax1, v_ax2])
+ T.writes(p_output0_intermediate[v_ax0, v_ax1, v_ax2])
+ p_output0_intermediate[v_ax0, v_ax1, v_ax2] = var_T_add_intermediate[v_ax0, v_ax1, v_ax2] + lv2062[v_ax0, v_ax1, v_ax2]
+
+ @T.prim_func
+ def func4(lv1387: T.Buffer((T.int64(320), T.int64(10240)), "uint32"), lv1388: T.Buffer((T.int64(80), T.int64(10240)), "float16"), lv2115: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16"), linear_bias196: T.Buffer((T.int64(10240),), "float32"), p_output0_intermediate: T.Buffer((T.int64(1), T.int64(1), T.int64(10240)), "float16")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ var_decode_intermediate = T.alloc_buffer((T.int64(2560), T.int64(10240)), "float16")
+ var_matmul_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(10240)))
+ var_T_add_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(10240)))
+ T_multiply = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(10240)))
+ compute = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(10240)))
+ T_multiply_1 = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(10240)))
+ T_add = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(10240)))
+ var_T_multiply_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(10240)))
+ for i, j in T.grid(T.int64(2560), T.int64(10240)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(lv1387[v_i // T.int64(8), v_j], lv1388[v_i // T.int64(32), v_j])
+ T.writes(var_decode_intermediate[v_i, v_j])
+ var_decode_intermediate[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv1387[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv1388[v_i // T.int64(32), v_j]
+ for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), T.int64(10240), T.int64(2560)):
+ with T.block("matmul"):
+ v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
+ T.reads(lv2115[v_i0, v_i1, v_k], var_decode_intermediate[v_k, v_i2])
+ T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
+ with T.init():
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float32(0)
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + T.Cast("float32", lv2115[v_i0, v_i1, v_k]) * T.Cast("float32", var_decode_intermediate[v_k, v_i2])
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(10240)):
+ with T.block("T_add"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_matmul_intermediate[v_ax0, v_ax1, v_ax2], linear_bias196[v_ax2])
+ T.writes(var_T_add_intermediate[v_ax0, v_ax1, v_ax2])
+ var_T_add_intermediate[v_ax0, v_ax1, v_ax2] = var_matmul_intermediate[v_ax0, v_ax1, v_ax2] + linear_bias196[v_ax2]
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(10240)):
+ with T.block("T_multiply"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_T_add_intermediate[v_ax0, v_ax1, v_ax2])
+ T.writes(T_multiply[v_ax0, v_ax1, v_ax2])
+ T_multiply[v_ax0, v_ax1, v_ax2] = var_T_add_intermediate[v_ax0, v_ax1, v_ax2] * T.float32(0.70710678118654757)
+ for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), T.int64(10240)):
+ with T.block("compute"):
+ v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(T_multiply[v_i0, v_i1, v_i2])
+ T.writes(compute[v_i0, v_i1, v_i2])
+ compute[v_i0, v_i1, v_i2] = T.erf(T_multiply[v_i0, v_i1, v_i2])
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(10240)):
+ with T.block("T_multiply_1"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(compute[v_ax0, v_ax1, v_ax2])
+ T.writes(T_multiply_1[v_ax0, v_ax1, v_ax2])
+ T_multiply_1[v_ax0, v_ax1, v_ax2] = compute[v_ax0, v_ax1, v_ax2] * T.float32(0.5)
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(10240)):
+ with T.block("T_add_1"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(T_multiply_1[v_ax0, v_ax1, v_ax2])
+ T.writes(T_add[v_ax0, v_ax1, v_ax2])
+ T_add[v_ax0, v_ax1, v_ax2] = T.float32(0.5) + T_multiply_1[v_ax0, v_ax1, v_ax2]
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(10240)):
+ with T.block("T_multiply_2"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_T_add_intermediate[v_ax0, v_ax1, v_ax2], T_add[v_ax0, v_ax1, v_ax2])
+ T.writes(var_T_multiply_intermediate[v_ax0, v_ax1, v_ax2])
+ var_T_multiply_intermediate[v_ax0, v_ax1, v_ax2] = var_T_add_intermediate[v_ax0, v_ax1, v_ax2] * T_add[v_ax0, v_ax1, v_ax2]
+ for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), T.int64(10240)):
+ with T.block("compute_1"):
+ v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(var_T_multiply_intermediate[v_i0, v_i1, v_i2])
+ T.writes(p_output0_intermediate[v_i0, v_i1, v_i2])
+ p_output0_intermediate[v_i0, v_i1, v_i2] = T.Cast("float16", var_T_multiply_intermediate[v_i0, v_i1, v_i2])
+
+ @T.prim_func
+ def func5(lv1393: T.Buffer((T.int64(1280), T.int64(2560)), "uint32"), lv1394: T.Buffer((T.int64(320), T.int64(2560)), "float16"), lv2121: T.Buffer((T.int64(1), T.int64(1), T.int64(10240)), "float16"), linear_bias197: T.Buffer((T.int64(2560),), "float32"), lv329: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16"), p_output0_intermediate: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ var_decode_intermediate = T.alloc_buffer((T.int64(10240), T.int64(2560)), "float16")
+ var_matmul_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)))
+ var_T_add_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)))
+ var_compute_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")
+ var_compute_intermediate_1 = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")
+ for i, j in T.grid(T.int64(10240), T.int64(2560)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(lv1393[v_i // T.int64(8), v_j], lv1394[v_i // T.int64(32), v_j])
+ T.writes(var_decode_intermediate[v_i, v_j])
+ var_decode_intermediate[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv1393[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv1394[v_i // T.int64(32), v_j]
+ for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), T.int64(2560), T.int64(10240)):
+ with T.block("matmul"):
+ v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
+ T.reads(lv2121[v_i0, v_i1, v_k], var_decode_intermediate[v_k, v_i2])
+ T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
+ with T.init():
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float32(0)
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + T.Cast("float32", lv2121[v_i0, v_i1, v_k]) * T.Cast("float32", var_decode_intermediate[v_k, v_i2])
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("T_add"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_matmul_intermediate[v_ax0, v_ax1, v_ax2], linear_bias197[v_ax2])
+ T.writes(var_T_add_intermediate[v_ax0, v_ax1, v_ax2])
+ var_T_add_intermediate[v_ax0, v_ax1, v_ax2] = var_matmul_intermediate[v_ax0, v_ax1, v_ax2] + linear_bias197[v_ax2]
+ for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("compute"):
+ v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(var_T_add_intermediate[v_i0, v_i1, v_i2])
+ T.writes(var_compute_intermediate[v_i0, v_i1, v_i2])
+ var_compute_intermediate[v_i0, v_i1, v_i2] = T.Cast("float16", var_T_add_intermediate[v_i0, v_i1, v_i2])
+ for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("compute_1"):
+ v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(var_compute_intermediate[v_i0, v_i1, v_i2])
+ T.writes(var_compute_intermediate_1[v_i0, v_i1, v_i2])
+ var_compute_intermediate_1[v_i0, v_i1, v_i2] = var_compute_intermediate[v_i0, v_i1, v_i2]
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("T_add_1"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_compute_intermediate_1[v_ax0, v_ax1, v_ax2], lv329[v_ax0, v_ax1, v_ax2])
+ T.writes(p_output0_intermediate[v_ax0, v_ax1, v_ax2])
+ p_output0_intermediate[v_ax0, v_ax1, v_ax2] = var_compute_intermediate_1[v_ax0, v_ax1, v_ax2] + lv329[v_ax0, v_ax1, v_ax2]
+
+ @T.prim_func
+ def func6(lv2509: T.Buffer((T.int64(1280), T.int64(2560)), "uint32"), lv2510: T.Buffer((T.int64(320), T.int64(2560)), "float16"), lv4105: T.Buffer((T.int64(1), T.int64(1), T.int64(10240)), "float16"), linear_bias383: T.Buffer((T.int64(2560),), "float32"), lv701: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16"), p_output0_intermediate: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float32")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ var_decode_intermediate = T.alloc_buffer((T.int64(10240), T.int64(2560)), "float16")
+ var_matmul_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)))
+ var_T_add_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)))
+ var_compute_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")
+ var_compute_intermediate_1 = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")
+ var_T_add_intermediate_1 = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")
+ for i, j in T.grid(T.int64(10240), T.int64(2560)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(lv2509[v_i // T.int64(8), v_j], lv2510[v_i // T.int64(32), v_j])
+ T.writes(var_decode_intermediate[v_i, v_j])
+ var_decode_intermediate[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv2509[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv2510[v_i // T.int64(32), v_j]
+ for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), T.int64(2560), T.int64(10240)):
+ with T.block("matmul"):
+ v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
+ T.reads(lv4105[v_i0, v_i1, v_k], var_decode_intermediate[v_k, v_i2])
+ T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
+ with T.init():
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float32(0)
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + T.Cast("float32", lv4105[v_i0, v_i1, v_k]) * T.Cast("float32", var_decode_intermediate[v_k, v_i2])
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("T_add"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_matmul_intermediate[v_ax0, v_ax1, v_ax2], linear_bias383[v_ax2])
+ T.writes(var_T_add_intermediate[v_ax0, v_ax1, v_ax2])
+ var_T_add_intermediate[v_ax0, v_ax1, v_ax2] = var_matmul_intermediate[v_ax0, v_ax1, v_ax2] + linear_bias383[v_ax2]
+ for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("compute"):
+ v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(var_T_add_intermediate[v_i0, v_i1, v_i2])
+ T.writes(var_compute_intermediate[v_i0, v_i1, v_i2])
+ var_compute_intermediate[v_i0, v_i1, v_i2] = T.Cast("float16", var_T_add_intermediate[v_i0, v_i1, v_i2])
+ for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("compute_1"):
+ v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(var_compute_intermediate[v_i0, v_i1, v_i2])
+ T.writes(var_compute_intermediate_1[v_i0, v_i1, v_i2])
+ var_compute_intermediate_1[v_i0, v_i1, v_i2] = var_compute_intermediate[v_i0, v_i1, v_i2]
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("T_add_1"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_compute_intermediate_1[v_ax0, v_ax1, v_ax2], lv701[v_ax0, v_ax1, v_ax2])
+ T.writes(var_T_add_intermediate_1[v_ax0, v_ax1, v_ax2])
+ var_T_add_intermediate_1[v_ax0, v_ax1, v_ax2] = var_compute_intermediate_1[v_ax0, v_ax1, v_ax2] + lv701[v_ax0, v_ax1, v_ax2]
+ for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("compute_2"):
+ v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
Review Comment:
will do
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on a diff in pull request #15141: [WIP] Scaffolding DLight
Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen commented on code in PR #15141:
URL: https://github.com/apache/tvm/pull/15141#discussion_r1239862854
##########
python/tvm/dlight/gpu/default_schedule.py:
##########
@@ -0,0 +1,89 @@
+# pylint: disable=missing-docstring
+from typing import Callable
+
+from tvm import tir
+from tvm._ffi import get_global_func
+from tvm.ir import IRModule
+from tvm.ir.transform import PassContext, module_pass
+from tvm.target import Target
+
+
+def _default_schedule( # pylint: disable=too-many-locals
+ func: tir.PrimFunc,
+ target: Target,
+) -> tir.PrimFunc:
+ max_threads_per_block = None
+ for name in ["max_threads_per_block", "max_num_threads"]:
+ if max_threads_per_block is None:
+ max_threads_per_block = target.attrs.get(name, None)
+ if max_threads_per_block is None:
+ max_threads_per_block = 64
+
+ get_loop_iter_type = get_global_func("tir.schedule.GetLoopIterType")
Review Comment:
This can be a python function get exposed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tvm-bot commented on pull request #15141: Scaffolding DLight
Posted by "tvm-bot (via GitHub)" <gi...@apache.org>.
tvm-bot commented on PR #15141:
URL: https://github.com/apache/tvm/pull/15141#issuecomment-1601993046
<!---bot-comment-->
Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from [Reviewers](https://github.com/apache/incubator-tvm/blob/master/CONTRIBUTORS.md#reviewers) by @-ing them in a comment.
<sub>Generated by [tvm-bot](https://github.com/apache/tvm/blob/main/ci/README.md#github-actions)</sub>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on a diff in pull request #15141: [WIP] Scaffolding DLight
Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen commented on code in PR #15141:
URL: https://github.com/apache/tvm/pull/15141#discussion_r1238660532
##########
main.py:
##########
@@ -0,0 +1,267 @@
+# pylint: disable=missing-docstring
+from typing import List, Optional
+
+from tvm import IRModule
+from tvm import meta_schedule as ms
+from tvm import tir
+from tvm.dlight import ScheduleGenerator, ScheduleRule, auto_inline_consumers
+
+
+class Decode(ScheduleRule):
+ def __init__(self):
+ ...
+
+ def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
+ pass
+
+ def clone(self) -> ScheduleRule:
+ return Decode()
+
+ def apply(self, mod: IRModule) -> Optional[tir.Schedule]: # pylint: disable=too-many-locals
+ sch = tir.Schedule(mod)
+ try:
+ decode = sch.get_block("decode")
+ except: # pylint: disable=bare-except
+ return None
+ len_tx: int = 8
+ len_ty: int = 8
+ len_yi: int = 1
+ len_yc: int = 8
+
+ # Step 1. Tile the decoding
+ i, j = sch.get_loops(decode)
+ by, ty, yi, yc = sch.split( # pylint: disable=invalid-name
+ i, factors=[None, len_ty, len_yi, len_yc]
+ )
+ bx, tx = sch.split(j, factors=[None, len_tx]) # pylint: disable=invalid-name
+ sch.reorder(by, bx, ty, tx, yi, yc)
+ sch.bind(by, "blockIdx.y")
+ sch.bind(bx, "blockIdx.x")
+ sch.bind(ty, "threadIdx.y")
+ sch.bind(tx, "threadIdx.x")
+ sch.unroll(yc)
+ # Step 2. Cache results in shared memory
+ rb = sch.cache_write(decode, 0, "shared") # pylint: disable=invalid-name
+ consumers = sch.get_consumers(rb)
+ if consumers:
+ (consumer,) = consumers
+ auto_inline_consumers(sch, consumer)
+ sch.compute_inline(rb)
+ rb = consumer # pylint: disable=invalid-name
+ # Step 3. Schedule the shared memory write back
+ sch.reverse_compute_at(rb, bx, preserve_unit_loops=True)
+ loop = sch.fuse(*sch.get_loops(rb)[-2:])
+ _, ty, tx, vec = sch.split( # pylint: disable=invalid-name
+ loop, factors=[None, len_ty, len_tx, 4]
+ )
+ sch.bind(ty, "threadIdx.y")
+ sch.bind(tx, "threadIdx.x")
+ sch.vectorize(vec)
+ sch.storage_align(decode, buffer_index=0, axis=0, factor=32, offset=1)
+ return sch
+
+
+class DecodeGemv(ScheduleRule):
+ def __init__(self):
+ ...
+
+ def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
+ pass
+
+ def clone(self) -> ScheduleRule:
+ return DecodeGemv()
+
+ def apply(self, mod: IRModule) -> Optional[tir.Schedule]: # pylint: disable=too-many-locals
+ sch = tir.Schedule(mod)
+ try:
+ gemv = sch.get_block("matmul")
+ decode = sch.get_block("decode")
+ except: # pylint: disable=bare-except
+ return None
+ len_vx: int = 2
+ len_tx: int = 64
+ len_km: int = 2
+ len_ki: int = 1 * 8
+ # Step 1. Schedule GEMV
+ # [b=1, i=1, j, k]
+ # split j => [b=1, i=1, (bx, vx, tx), k]
+ # fuse (b, i, bx) => [bx, vx, tx, (k)]
+ # split k => [bx, vx, tx, (ko, k_m, ki * 8)]
+ rb = sch.cache_write(gemv, 0, "local") # pylint: disable=invalid-name
+ b, i, j, k = sch.get_loops(gemv) # pylint: disable=invalid-name
+ assert sch.get(b).extent.value == 1
+ assert sch.get(i).extent.value == 1
+ bx, vx, tx = sch.split(j, [None, len_vx, len_tx]) # pylint: disable=invalid-name
+ bx = sch.fuse(b, i, bx) # pylint: disable=invalid-name
+ k_o, k_m, k_i = sch.split(k, [None, len_km, len_ki])
+ sch.bind(bx, thread_axis="blockIdx.x")
+ sch.bind(vx, thread_axis="vthread.x")
+ sch.bind(tx, thread_axis="threadIdx.x")
+ sch.reorder(bx, vx, tx, k_o, k_m, k_i)
+ sch.unroll(k_i)
+ # Step 2. Schedule decode: move to under threadIdx.x and fetch separately for each thread
+ sch.compute_at(decode, k_m, preserve_unit_loops=True)
+ sch.set_scope(decode, 0, "local")
+ _, unroll = sch.split(sch.get_loops(decode)[-2], [None, 8])
+ sch.unroll(unroll)
+
+ # Step 3. Cooperative fetch GEMV
+ def cooperative_fetch(block, tx): # pylint: disable=invalid-name
+ block = sch.cache_read(block, 0, "shared")
+ sch.compute_at(block, tx, preserve_unit_loops=True)
+ loop = sch.fuse(*sch.get_loops(block)[-2:])
+ len_vector = sch.sample_categorical(
+ [1, 2, 3, 4],
+ probs=[0.25, 0.25, 0.25, 0.25],
+ )
+ _, tx, vec = sch.split(loop, [None, len_tx, len_vector])
+ sch.bind(tx, thread_axis="threadIdx.x")
+ sch.vectorize(vec)
+ sch.storage_align(block, buffer_index=0, axis=-2, factor=32, offset=8)
+
+ cooperative_fetch(gemv, k_o)
+ # Step 4. Schedule epilogue
+ auto_inline_consumers(sch, rb)
+ sch.reverse_compute_at(rb, tx, preserve_unit_loops=True)
+ # Step 5. Postprocess: decompose reduction
+ sch.decompose_reduction(gemv, k_o)
+ return [sch]
+
+
+class Normalization(ScheduleRule):
+ def __init__(self):
+ ...
+
+ def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
+ pass
+
+ def clone(self) -> ScheduleRule:
+ return Normalization()
+
+ def apply(self, mod: IRModule) -> Optional[tir.Schedule]: # pylint: disable=too-many-locals
+ sch = tir.Schedule(mod)
+ b_reduce = None
+ for name in ["Ared_temp", "A_red_temp"]:
+ try:
+ b_reduce = sch.get_block(name)
+ except:
+ continue
+ else:
+ break
+ if b_reduce is None:
+ return None
+ len_tx: int = 256
+ unroll_depth: int = 256
+
+ (b_spatial,) = sch.get_consumers(b_reduce)
+ loops = sch.get_loops(b_spatial)
+ bx = sch.fuse(*loops[:-1]) # pylint: disable=invalid-name
+ _, tx = sch.split(loops[-1], [None, len_tx]) # pylint: disable=invalid-name
+ sch.bind(bx, "blockIdx.x")
+ sch.bind(tx, "threadIdx.x")
+
+ for i, _ in enumerate(sch.get(b_reduce).writes):
+ sch.set_scope(b_reduce, buffer_index=i, storage_scope="shared")
+ sch.compute_at(b_reduce, bx, preserve_unit_loops=True)
+ _, tx = sch.split( # pylint: disable=invalid-name
+ sch.get_loops(b_reduce)[-1],
+ [None, len_tx],
+ )
+ sch.bind(tx, "threadIdx.x")
+ auto_inline_consumers(sch, b_spatial)
+ sch.annotate(bx, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
+ sch.annotate(bx, ann_key="pragma_unroll_explicit", ann_val=1)
+
+ return sch
+
+
+class Softmax(ScheduleRule):
+ def __init__(self):
+ ...
+
+ def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
+ pass
+
+ def clone(self) -> ScheduleRule:
+ return Softmax()
+
+ def apply(self, mod: IRModule) -> Optional[tir.Schedule]: # pylint: disable=too-many-locals
+ sch = tir.Schedule(mod)
+ try:
+ b_reduce_0 = sch.get_block("T_softmax_maxelem")
+ b_reduce_1 = sch.get_block("T_softmax_expsum")
+ b_spatial = sch.get_block("T_softmax_norm")
+ except: # pylint: disable=bare-except
+ return None
+
+ len_tx: int = 256
+ unroll_depth: int = 256
+
+ sch.compute_inline(sch.get_producers(b_reduce_1)[0])
+
+ loops = sch.get_loops(b_spatial)
+ bx = sch.fuse(*loops[:-1]) # pylint: disable=invalid-name
+ _, tx = sch.split(loops[-1], [None, len_tx]) # pylint: disable=invalid-name
+ sch.bind(bx, "blockIdx.x")
+ sch.bind(tx, "threadIdx.x")
+
+ sch.set_scope(b_reduce_1, buffer_index=0, storage_scope="shared")
+ sch.compute_at(b_reduce_1, bx, preserve_unit_loops=True)
+ _, tx = sch.split( # pylint: disable=invalid-name
+ sch.get_loops(b_reduce_1)[-1],
+ [None, len_tx],
+ )
+ sch.bind(tx, "threadIdx.x")
+
+ sch.set_scope(b_reduce_0, buffer_index=0, storage_scope="shared")
+ sch.compute_at(b_reduce_0, bx, preserve_unit_loops=True)
+ _, tx = sch.split( # pylint: disable=invalid-name
+ sch.get_loops(b_reduce_0)[-1],
+ [None, len_tx],
+ )
+ sch.bind(tx, "threadIdx.x")
+
+ sch.annotate(bx, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
+ sch.annotate(bx, ann_key="pragma_unroll_explicit", ann_val=1)
+
+ auto_inline_consumers(sch, b_spatial)
+ return [sch]
+
+
+def main():
+ from tvm.dlight.testing import mod_decode, mod_decode_gemv, mod_norm, mod_softmax
+
+ gen = ScheduleGenerator(
+ rules=[
+ DecodeGemv(),
+ Decode(),
+ Normalization(),
+ Softmax(),
+ ]
+ )
+
+ for py_mod in [
+ mod_decode,
+ mod_decode_gemv,
+ mod_norm,
+ mod_softmax, # Needs to upstream `compute-inline`
+ ]:
+ i = 1
+ while True:
+ try:
+ func = py_mod.Module[f"func{i}"]
+ except: # pylint: disable=bare-except
+ break
+ else:
+ print(f"Working on {py_mod}::func{i}")
+ i += 1
+ mod = IRModule.from_expr(func.with_attr("global_symbol", "main"))
+ schedules = gen.generate_design_space(mod)
Review Comment:
Let us structure the testcase UT same as Before/After test, e.g. one softmax function before and after with structural equality test.
Likely most of the modules in dlight testing can move to `tests/dlight/` (that can be added to CI)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #15141: [WIP] Scaffolding DLight
Posted by "junrushao (via GitHub)" <gi...@apache.org>.
junrushao commented on code in PR #15141:
URL: https://github.com/apache/tvm/pull/15141#discussion_r1239919727
##########
python/tvm/dlight/gpu/default_schedule.py:
##########
@@ -0,0 +1,89 @@
+# pylint: disable=missing-docstring
+from typing import Callable
+
+from tvm import tir
+from tvm._ffi import get_global_func
+from tvm.ir import IRModule
+from tvm.ir.transform import PassContext, module_pass
+from tvm.target import Target
+
+
+def _default_schedule( # pylint: disable=too-many-locals
+ func: tir.PrimFunc,
+ target: Target,
+) -> tir.PrimFunc:
+ max_threads_per_block = None
+ for name in ["max_threads_per_block", "max_num_threads"]:
+ if max_threads_per_block is None:
+ max_threads_per_block = target.attrs.get(name, None)
+ if max_threads_per_block is None:
Review Comment:
that makes sense
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on a diff in pull request #15141: [WIP] Scaffolding DLight
Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen commented on code in PR #15141:
URL: https://github.com/apache/tvm/pull/15141#discussion_r1238660532
##########
main.py:
##########
@@ -0,0 +1,267 @@
+# pylint: disable=missing-docstring
+from typing import List, Optional
+
+from tvm import IRModule
+from tvm import meta_schedule as ms
+from tvm import tir
+from tvm.dlight import ScheduleGenerator, ScheduleRule, auto_inline_consumers
+
+
+class Decode(ScheduleRule):
+ def __init__(self):
+ ...
+
+ def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
+ pass
+
+ def clone(self) -> ScheduleRule:
+ return Decode()
+
+ def apply(self, mod: IRModule) -> Optional[tir.Schedule]: # pylint: disable=too-many-locals
+ sch = tir.Schedule(mod)
+ try:
+ decode = sch.get_block("decode")
+ except: # pylint: disable=bare-except
+ return None
+ len_tx: int = 8
+ len_ty: int = 8
+ len_yi: int = 1
+ len_yc: int = 8
+
+ # Step 1. Tile the decoding
+ i, j = sch.get_loops(decode)
+ by, ty, yi, yc = sch.split( # pylint: disable=invalid-name
+ i, factors=[None, len_ty, len_yi, len_yc]
+ )
+ bx, tx = sch.split(j, factors=[None, len_tx]) # pylint: disable=invalid-name
+ sch.reorder(by, bx, ty, tx, yi, yc)
+ sch.bind(by, "blockIdx.y")
+ sch.bind(bx, "blockIdx.x")
+ sch.bind(ty, "threadIdx.y")
+ sch.bind(tx, "threadIdx.x")
+ sch.unroll(yc)
+ # Step 2. Cache results in shared memory
+ rb = sch.cache_write(decode, 0, "shared") # pylint: disable=invalid-name
+ consumers = sch.get_consumers(rb)
+ if consumers:
+ (consumer,) = consumers
+ auto_inline_consumers(sch, consumer)
+ sch.compute_inline(rb)
+ rb = consumer # pylint: disable=invalid-name
+ # Step 3. Schedule the shared memory write back
+ sch.reverse_compute_at(rb, bx, preserve_unit_loops=True)
+ loop = sch.fuse(*sch.get_loops(rb)[-2:])
+ _, ty, tx, vec = sch.split( # pylint: disable=invalid-name
+ loop, factors=[None, len_ty, len_tx, 4]
+ )
+ sch.bind(ty, "threadIdx.y")
+ sch.bind(tx, "threadIdx.x")
+ sch.vectorize(vec)
+ sch.storage_align(decode, buffer_index=0, axis=0, factor=32, offset=1)
+ return sch
+
+
+class DecodeGemv(ScheduleRule):
+ def __init__(self):
+ ...
+
+ def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
+ pass
+
+ def clone(self) -> ScheduleRule:
+ return DecodeGemv()
+
+ def apply(self, mod: IRModule) -> Optional[tir.Schedule]: # pylint: disable=too-many-locals
+ sch = tir.Schedule(mod)
+ try:
+ gemv = sch.get_block("matmul")
+ decode = sch.get_block("decode")
+ except: # pylint: disable=bare-except
+ return None
+ len_vx: int = 2
+ len_tx: int = 64
+ len_km: int = 2
+ len_ki: int = 1 * 8
+ # Step 1. Schedule GEMV
+ # [b=1, i=1, j, k]
+ # split j => [b=1, i=1, (bx, vx, tx), k]
+ # fuse (b, i, bx) => [bx, vx, tx, (k)]
+ # split k => [bx, vx, tx, (ko, k_m, ki * 8)]
+ rb = sch.cache_write(gemv, 0, "local") # pylint: disable=invalid-name
+ b, i, j, k = sch.get_loops(gemv) # pylint: disable=invalid-name
+ assert sch.get(b).extent.value == 1
+ assert sch.get(i).extent.value == 1
+ bx, vx, tx = sch.split(j, [None, len_vx, len_tx]) # pylint: disable=invalid-name
+ bx = sch.fuse(b, i, bx) # pylint: disable=invalid-name
+ k_o, k_m, k_i = sch.split(k, [None, len_km, len_ki])
+ sch.bind(bx, thread_axis="blockIdx.x")
+ sch.bind(vx, thread_axis="vthread.x")
+ sch.bind(tx, thread_axis="threadIdx.x")
+ sch.reorder(bx, vx, tx, k_o, k_m, k_i)
+ sch.unroll(k_i)
+ # Step 2. Schedule decode: move to under threadIdx.x and fetch separately for each thread
+ sch.compute_at(decode, k_m, preserve_unit_loops=True)
+ sch.set_scope(decode, 0, "local")
+ _, unroll = sch.split(sch.get_loops(decode)[-2], [None, 8])
+ sch.unroll(unroll)
+
+ # Step 3. Cooperative fetch GEMV
+ def cooperative_fetch(block, tx): # pylint: disable=invalid-name
+ block = sch.cache_read(block, 0, "shared")
+ sch.compute_at(block, tx, preserve_unit_loops=True)
+ loop = sch.fuse(*sch.get_loops(block)[-2:])
+ len_vector = sch.sample_categorical(
+ [1, 2, 3, 4],
+ probs=[0.25, 0.25, 0.25, 0.25],
+ )
+ _, tx, vec = sch.split(loop, [None, len_tx, len_vector])
+ sch.bind(tx, thread_axis="threadIdx.x")
+ sch.vectorize(vec)
+ sch.storage_align(block, buffer_index=0, axis=-2, factor=32, offset=8)
+
+ cooperative_fetch(gemv, k_o)
+ # Step 4. Schedule epilogue
+ auto_inline_consumers(sch, rb)
+ sch.reverse_compute_at(rb, tx, preserve_unit_loops=True)
+ # Step 5. Postprocess: decompose reduction
+ sch.decompose_reduction(gemv, k_o)
+ return [sch]
+
+
+class Normalization(ScheduleRule):
+ def __init__(self):
+ ...
+
+ def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
+ pass
+
+ def clone(self) -> ScheduleRule:
+ return Normalization()
+
+ def apply(self, mod: IRModule) -> Optional[tir.Schedule]: # pylint: disable=too-many-locals
+ sch = tir.Schedule(mod)
+ b_reduce = None
+ for name in ["Ared_temp", "A_red_temp"]:
+ try:
+ b_reduce = sch.get_block(name)
+ except:
+ continue
+ else:
+ break
+ if b_reduce is None:
+ return None
+ len_tx: int = 256
+ unroll_depth: int = 256
+
+ (b_spatial,) = sch.get_consumers(b_reduce)
+ loops = sch.get_loops(b_spatial)
+ bx = sch.fuse(*loops[:-1]) # pylint: disable=invalid-name
+ _, tx = sch.split(loops[-1], [None, len_tx]) # pylint: disable=invalid-name
+ sch.bind(bx, "blockIdx.x")
+ sch.bind(tx, "threadIdx.x")
+
+ for i, _ in enumerate(sch.get(b_reduce).writes):
+ sch.set_scope(b_reduce, buffer_index=i, storage_scope="shared")
+ sch.compute_at(b_reduce, bx, preserve_unit_loops=True)
+ _, tx = sch.split( # pylint: disable=invalid-name
+ sch.get_loops(b_reduce)[-1],
+ [None, len_tx],
+ )
+ sch.bind(tx, "threadIdx.x")
+ auto_inline_consumers(sch, b_spatial)
+ sch.annotate(bx, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
+ sch.annotate(bx, ann_key="pragma_unroll_explicit", ann_val=1)
+
+ return sch
+
+
+class Softmax(ScheduleRule):
+ def __init__(self):
+ ...
+
+ def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
+ pass
+
+ def clone(self) -> ScheduleRule:
+ return Softmax()
+
+ def apply(self, mod: IRModule) -> Optional[tir.Schedule]: # pylint: disable=too-many-locals
+ sch = tir.Schedule(mod)
+ try:
+ b_reduce_0 = sch.get_block("T_softmax_maxelem")
+ b_reduce_1 = sch.get_block("T_softmax_expsum")
+ b_spatial = sch.get_block("T_softmax_norm")
+ except: # pylint: disable=bare-except
+ return None
+
+ len_tx: int = 256
+ unroll_depth: int = 256
+
+ sch.compute_inline(sch.get_producers(b_reduce_1)[0])
+
+ loops = sch.get_loops(b_spatial)
+ bx = sch.fuse(*loops[:-1]) # pylint: disable=invalid-name
+ _, tx = sch.split(loops[-1], [None, len_tx]) # pylint: disable=invalid-name
+ sch.bind(bx, "blockIdx.x")
+ sch.bind(tx, "threadIdx.x")
+
+ sch.set_scope(b_reduce_1, buffer_index=0, storage_scope="shared")
+ sch.compute_at(b_reduce_1, bx, preserve_unit_loops=True)
+ _, tx = sch.split( # pylint: disable=invalid-name
+ sch.get_loops(b_reduce_1)[-1],
+ [None, len_tx],
+ )
+ sch.bind(tx, "threadIdx.x")
+
+ sch.set_scope(b_reduce_0, buffer_index=0, storage_scope="shared")
+ sch.compute_at(b_reduce_0, bx, preserve_unit_loops=True)
+ _, tx = sch.split( # pylint: disable=invalid-name
+ sch.get_loops(b_reduce_0)[-1],
+ [None, len_tx],
+ )
+ sch.bind(tx, "threadIdx.x")
+
+ sch.annotate(bx, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
+ sch.annotate(bx, ann_key="pragma_unroll_explicit", ann_val=1)
+
+ auto_inline_consumers(sch, b_spatial)
+ return [sch]
+
+
+def main():
+ from tvm.dlight.testing import mod_decode, mod_decode_gemv, mod_norm, mod_softmax
+
+ gen = ScheduleGenerator(
+ rules=[
+ DecodeGemv(),
+ Decode(),
+ Normalization(),
+ Softmax(),
+ ]
+ )
+
+ for py_mod in [
+ mod_decode,
+ mod_decode_gemv,
+ mod_norm,
+ mod_softmax, # Needs to upstream `compute-inline`
+ ]:
+ i = 1
+ while True:
+ try:
+ func = py_mod.Module[f"func{i}"]
+ except: # pylint: disable=bare-except
+ break
+ else:
+ print(f"Working on {py_mod}::func{i}")
+ i += 1
+ mod = IRModule.from_expr(func.with_attr("global_symbol", "main"))
+ schedules = gen.generate_design_space(mod)
Review Comment:
Let us structure the testcase UT same as Before/After test, e.g. one softmax function before and after with structural equality test. Likely most of the dlight testing can move to `tests/dlight/` (that can be added to CI)
Atm we only test that the schedule exists
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen merged pull request #15141: [Unity] Scaffolding DLight
Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen merged PR #15141:
URL: https://github.com/apache/tvm/pull/15141
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #15141: [WIP] Scaffolding DLight
Posted by "junrushao (via GitHub)" <gi...@apache.org>.
junrushao commented on code in PR #15141:
URL: https://github.com/apache/tvm/pull/15141#discussion_r1239919253
##########
python/tvm/dlight/testing/modules.py:
##########
@@ -0,0 +1,2621 @@
+from tvm.script import ir as I
+from tvm.script import tir as T
+
+# fmt: off
+
+@I.ir_module
+class Decode:
+
Review Comment:
got it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on a diff in pull request #15141: [WIP] Scaffolding DLight
Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen commented on code in PR #15141:
URL: https://github.com/apache/tvm/pull/15141#discussion_r1239861339
##########
python/tvm/dlight/testing/modules.py:
##########
@@ -0,0 +1,2621 @@
+from tvm.script import ir as I
+from tvm.script import tir as T
+
+# fmt: off
+
+@I.ir_module
+class Decode:
+
Review Comment:
let us not have testing here, instead move to `tests/python/dlight/` so we can grow test cases organically
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #15141: [WIP] Scaffolding DLight
Posted by "junrushao (via GitHub)" <gi...@apache.org>.
junrushao commented on code in PR #15141:
URL: https://github.com/apache/tvm/pull/15141#discussion_r1239922815
##########
python/tvm/dlight/gpu/default_schedule.py:
##########
@@ -0,0 +1,89 @@
+# pylint: disable=missing-docstring
+from typing import Callable
+
+from tvm import tir
+from tvm._ffi import get_global_func
+from tvm.ir import IRModule
+from tvm.ir.transform import PassContext, module_pass
+from tvm.target import Target
+
+
+def _default_schedule( # pylint: disable=too-many-locals
+ func: tir.PrimFunc,
+ target: Target,
+) -> tir.PrimFunc:
+ max_threads_per_block = None
+ for name in ["max_threads_per_block", "max_num_threads"]:
+ if max_threads_per_block is None:
+ max_threads_per_block = target.attrs.get(name, None)
+ if max_threads_per_block is None:
Review Comment:
i'm going to remove this rule first. seems not super relevant to this PR
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on a diff in pull request #15141: [WIP] Scaffolding DLight
Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen commented on code in PR #15141:
URL: https://github.com/apache/tvm/pull/15141#discussion_r1239862522
##########
python/tvm/dlight/gpu/default_schedule.py:
##########
@@ -0,0 +1,89 @@
+# pylint: disable=missing-docstring
+from typing import Callable
+
+from tvm import tir
+from tvm._ffi import get_global_func
+from tvm.ir import IRModule
+from tvm.ir.transform import PassContext, module_pass
+from tvm.target import Target
+
+
+def _default_schedule( # pylint: disable=too-many-locals
+ func: tir.PrimFunc,
+ target: Target,
+) -> tir.PrimFunc:
+ max_threads_per_block = None
+ for name in ["max_threads_per_block", "max_num_threads"]:
+ if max_threads_per_block is None:
+ max_threads_per_block = target.attrs.get(name, None)
+ if max_threads_per_block is None:
Review Comment:
Should have one schedule_rule, and invoke that schedule rule here. The schedule do not need to match all instances. Dlight is meant to get performant schedules
##########
python/tvm/dlight/gpu/default_schedule.py:
##########
@@ -0,0 +1,89 @@
+# pylint: disable=missing-docstring
+from typing import Callable
+
+from tvm import tir
+from tvm._ffi import get_global_func
+from tvm.ir import IRModule
+from tvm.ir.transform import PassContext, module_pass
+from tvm.target import Target
+
+
+def _default_schedule( # pylint: disable=too-many-locals
+ func: tir.PrimFunc,
+ target: Target,
+) -> tir.PrimFunc:
+ max_threads_per_block = None
+ for name in ["max_threads_per_block", "max_num_threads"]:
+ if max_threads_per_block is None:
+ max_threads_per_block = target.attrs.get(name, None)
+ if max_threads_per_block is None:
Review Comment:
Should have one schedule_rule, and invoke that schedule rule here. The schedule do not need to match all instances (match tag is fine). Dlight is meant to get performant schedules
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on a diff in pull request #15141: [WIP] Scaffolding DLight
Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen commented on code in PR #15141:
URL: https://github.com/apache/tvm/pull/15141#discussion_r1240021306
##########
tests/python/dlight/test_modules.py:
##########
@@ -0,0 +1,2621 @@
+from tvm.script import ir as I
+from tvm.script import tir as T
+
+# fmt: off
+
+@I.ir_module
+class Decode:
+
+ @T.prim_func
+ def func1(A: T.Buffer((T.int64(320), T.int64(2560)), "uint32"), B: T.Buffer((T.int64(80), T.int64(2560)), "float16"), T_transpose: T.Buffer((T.int64(2560), T.int64(2560)), "float16")):
+ T.func_attr({"op_pattern": 2, "tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ decode = T.alloc_buffer((T.int64(2560), T.int64(2560)), "float16")
+ for i, j in T.grid(T.int64(2560), T.int64(2560)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(A[v_i // T.int64(8), v_j], B[v_i // T.int64(32), v_j])
+ T.writes(decode[v_i, v_j])
+ decode[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(A[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * B[v_i // T.int64(32), v_j]
+ for ax0, ax1 in T.grid(T.int64(2560), T.int64(2560)):
+ with T.block("T_transpose"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(decode[v_ax1, v_ax0])
+ T.writes(T_transpose[v_ax0, v_ax1])
+ T_transpose[v_ax0, v_ax1] = decode[v_ax1, v_ax0]
+
+ @T.prim_func
+ def func2(A: T.Buffer((T.int64(320), T.int64(10240)), "uint32"), B: T.Buffer((T.int64(80), T.int64(10240)), "float16"), T_transpose: T.Buffer((T.int64(10240), T.int64(2560)), "float16")):
+ T.func_attr({"op_pattern": 2, "tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ decode = T.alloc_buffer((T.int64(2560), T.int64(10240)), "float16")
+ for i, j in T.grid(T.int64(2560), T.int64(10240)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(A[v_i // T.int64(8), v_j], B[v_i // T.int64(32), v_j])
+ T.writes(decode[v_i, v_j])
+ decode[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(A[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * B[v_i // T.int64(32), v_j]
+ for ax0, ax1 in T.grid(T.int64(10240), T.int64(2560)):
+ with T.block("T_transpose"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(decode[v_ax1, v_ax0])
+ T.writes(T_transpose[v_ax0, v_ax1])
+ T_transpose[v_ax0, v_ax1] = decode[v_ax1, v_ax0]
+
+ @T.prim_func
+ def func3(A: T.Buffer((T.int64(1280), T.int64(2560)), "uint32"), B: T.Buffer((T.int64(320), T.int64(2560)), "float16"), T_transpose: T.Buffer((T.int64(2560), T.int64(10240)), "float16")):
+ T.func_attr({"op_pattern": 2, "tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ decode = T.alloc_buffer((T.int64(10240), T.int64(2560)), "float16")
+ for i, j in T.grid(T.int64(10240), T.int64(2560)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(A[v_i // T.int64(8), v_j], B[v_i // T.int64(32), v_j])
+ T.writes(decode[v_i, v_j])
+ decode[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(A[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * B[v_i // T.int64(32), v_j]
+ for ax0, ax1 in T.grid(T.int64(2560), T.int64(10240)):
+ with T.block("T_transpose"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(decode[v_ax1, v_ax0])
+ T.writes(T_transpose[v_ax0, v_ax1])
+ T_transpose[v_ax0, v_ax1] = decode[v_ax1, v_ax0]
+
+ @T.prim_func
+ def func4(A: T.Buffer((T.int64(512), T.int64(4096)), "uint32"), B: T.Buffer((T.int64(128), T.int64(4096)), "float16"), T_transpose: T.Buffer((T.int64(4096), T.int64(4096)), "float16")):
+ T.func_attr({"op_pattern": 2, "tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ decode = T.alloc_buffer((T.int64(4096), T.int64(4096)), "float16")
+ for i, j in T.grid(T.int64(4096), T.int64(4096)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(A[v_i // T.int64(8), v_j], B[v_i // T.int64(32), v_j])
+ T.writes(decode[v_i, v_j])
+ decode[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(A[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * B[v_i // T.int64(32), v_j]
+ for ax0, ax1 in T.grid(T.int64(4096), T.int64(4096)):
+ with T.block("T_transpose"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(decode[v_ax1, v_ax0])
+ T.writes(T_transpose[v_ax0, v_ax1])
+ T_transpose[v_ax0, v_ax1] = decode[v_ax1, v_ax0]
+
+ @T.prim_func
+ def func5(A: T.Buffer((T.int64(512), T.int64(11008)), "uint32"), B: T.Buffer((T.int64(128), T.int64(11008)), "float16"), T_transpose: T.Buffer((T.int64(11008), T.int64(4096)), "float16")):
+ T.func_attr({"op_pattern": 2, "tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ decode = T.alloc_buffer((T.int64(4096), T.int64(11008)), "float16")
+ for i, j in T.grid(T.int64(4096), T.int64(11008)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(A[v_i // T.int64(8), v_j], B[v_i // T.int64(32), v_j])
+ T.writes(decode[v_i, v_j])
+ decode[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(A[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * B[v_i // T.int64(32), v_j]
+ for ax0, ax1 in T.grid(T.int64(11008), T.int64(4096)):
+ with T.block("T_transpose"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(decode[v_ax1, v_ax0])
+ T.writes(T_transpose[v_ax0, v_ax1])
+ T_transpose[v_ax0, v_ax1] = decode[v_ax1, v_ax0]
+
+ @T.prim_func
+ def func6(A: T.Buffer((T.int64(1376), T.int64(4096)), "uint32"), B: T.Buffer((T.int64(344), T.int64(4096)), "float16"), T_transpose: T.Buffer((T.int64(4096), T.int64(11008)), "float16")):
+ T.func_attr({"op_pattern": 2, "tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ decode = T.alloc_buffer((T.int64(11008), T.int64(4096)), "float16")
+ for i, j in T.grid(T.int64(11008), T.int64(4096)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(A[v_i // T.int64(8), v_j], B[v_i // T.int64(32), v_j])
+ T.writes(decode[v_i, v_j])
+ decode[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(A[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * B[v_i // T.int64(32), v_j]
+ for ax0, ax1 in T.grid(T.int64(4096), T.int64(11008)):
+ with T.block("T_transpose"):
+ v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
+ T.reads(decode[v_ax1, v_ax0])
+ T.writes(T_transpose[v_ax0, v_ax1])
+ T_transpose[v_ax0, v_ax1] = decode[v_ax1, v_ax0]
+
+
+@I.ir_module
+class DecodeGemv:
+
+ @T.prim_func
+ def func1(lv2515: T.Buffer((T.int64(320), T.int64(50432)), "uint32"), lv2516: T.Buffer((T.int64(80), T.int64(50432)), "float32"), lv705: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float32"), var_matmul_intermediate: T.Buffer((T.int64(1), T.int64(1), T.int64(50432)), "float32")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ var_decode_intermediate = T.alloc_buffer((T.int64(2560), T.int64(50432)))
+ for i, j in T.grid(T.int64(2560), T.int64(50432)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(lv2515[v_i // T.int64(8), v_j], lv2516[v_i // T.int64(32), v_j])
+ T.writes(var_decode_intermediate[v_i, v_j])
+ var_decode_intermediate[v_i, v_j] = T.Cast("float32", T.Cast("float16", T.bitwise_and(T.shift_right(lv2515[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv2516[v_i // T.int64(32), v_j]
+ for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), T.int64(50432), T.int64(2560)):
+ with T.block("matmul"):
+ v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
+ T.reads(lv705[v_i0, v_i1, v_k], var_decode_intermediate[v_k, v_i2])
+ T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
+ with T.init():
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float32(0)
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + lv705[v_i0, v_i1, v_k] * var_decode_intermediate[v_k, v_i2]
+
+ @T.prim_func
+ def func2(lv1363: T.Buffer((T.int64(320), T.int64(2560)), "uint32"), lv1364: T.Buffer((T.int64(80), T.int64(2560)), "float16"), lv2067: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16"), linear_bias192: T.Buffer((T.int64(2560),), "float16"), p_output0_intermediate: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ var_decode_intermediate = T.alloc_buffer((T.int64(2560), T.int64(2560)), "float16")
+ var_matmul_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")
+ for i, j in T.grid(T.int64(2560), T.int64(2560)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(lv1363[v_i // T.int64(8), v_j], lv1364[v_i // T.int64(32), v_j])
+ T.writes(var_decode_intermediate[v_i, v_j])
+ var_decode_intermediate[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv1363[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv1364[v_i // T.int64(32), v_j]
+ for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), T.int64(2560), T.int64(2560)):
+ with T.block("matmul"):
+ v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
+ T.reads(lv2067[v_i0, v_i1, v_k], var_decode_intermediate[v_k, v_i2])
+ T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
+ with T.init():
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0)
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + lv2067[v_i0, v_i1, v_k] * var_decode_intermediate[v_k, v_i2]
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("T_add"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_matmul_intermediate[v_ax0, v_ax1, v_ax2], linear_bias192[v_ax2])
+ T.writes(p_output0_intermediate[v_ax0, v_ax1, v_ax2])
+ p_output0_intermediate[v_ax0, v_ax1, v_ax2] = var_matmul_intermediate[v_ax0, v_ax1, v_ax2] + linear_bias192[v_ax2]
+
+ @T.prim_func
+ def func3(lv1381: T.Buffer((T.int64(320), T.int64(2560)), "uint32"), lv1382: T.Buffer((T.int64(80), T.int64(2560)), "float16"), lv328: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16"), linear_bias195: T.Buffer((T.int64(2560),), "float16"), lv2062: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16"), p_output0_intermediate: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ var_decode_intermediate = T.alloc_buffer((T.int64(2560), T.int64(2560)), "float16")
+ var_matmul_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")
+ var_T_add_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")
+ for i, j in T.grid(T.int64(2560), T.int64(2560)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(lv1381[v_i // T.int64(8), v_j], lv1382[v_i // T.int64(32), v_j])
+ T.writes(var_decode_intermediate[v_i, v_j])
+ var_decode_intermediate[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv1381[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv1382[v_i // T.int64(32), v_j]
+ for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), T.int64(2560), T.int64(2560)):
+ with T.block("matmul"):
+ v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
+ T.reads(lv328[v_i0, v_i1, v_k], var_decode_intermediate[v_k, v_i2])
+ T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
+ with T.init():
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float16(0)
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + lv328[v_i0, v_i1, v_k] * var_decode_intermediate[v_k, v_i2]
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("T_add"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_matmul_intermediate[v_ax0, v_ax1, v_ax2], linear_bias195[v_ax2])
+ T.writes(var_T_add_intermediate[v_ax0, v_ax1, v_ax2])
+ var_T_add_intermediate[v_ax0, v_ax1, v_ax2] = var_matmul_intermediate[v_ax0, v_ax1, v_ax2] + linear_bias195[v_ax2]
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("T_add_1"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_T_add_intermediate[v_ax0, v_ax1, v_ax2], lv2062[v_ax0, v_ax1, v_ax2])
+ T.writes(p_output0_intermediate[v_ax0, v_ax1, v_ax2])
+ p_output0_intermediate[v_ax0, v_ax1, v_ax2] = var_T_add_intermediate[v_ax0, v_ax1, v_ax2] + lv2062[v_ax0, v_ax1, v_ax2]
+
+ @T.prim_func
+ def func4(lv1387: T.Buffer((T.int64(320), T.int64(10240)), "uint32"), lv1388: T.Buffer((T.int64(80), T.int64(10240)), "float16"), lv2115: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16"), linear_bias196: T.Buffer((T.int64(10240),), "float32"), p_output0_intermediate: T.Buffer((T.int64(1), T.int64(1), T.int64(10240)), "float16")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ var_decode_intermediate = T.alloc_buffer((T.int64(2560), T.int64(10240)), "float16")
+ var_matmul_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(10240)))
+ var_T_add_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(10240)))
+ T_multiply = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(10240)))
+ compute = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(10240)))
+ T_multiply_1 = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(10240)))
+ T_add = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(10240)))
+ var_T_multiply_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(10240)))
+ for i, j in T.grid(T.int64(2560), T.int64(10240)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(lv1387[v_i // T.int64(8), v_j], lv1388[v_i // T.int64(32), v_j])
+ T.writes(var_decode_intermediate[v_i, v_j])
+ var_decode_intermediate[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv1387[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv1388[v_i // T.int64(32), v_j]
+ for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), T.int64(10240), T.int64(2560)):
+ with T.block("matmul"):
+ v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
+ T.reads(lv2115[v_i0, v_i1, v_k], var_decode_intermediate[v_k, v_i2])
+ T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
+ with T.init():
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float32(0)
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + T.Cast("float32", lv2115[v_i0, v_i1, v_k]) * T.Cast("float32", var_decode_intermediate[v_k, v_i2])
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(10240)):
+ with T.block("T_add"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_matmul_intermediate[v_ax0, v_ax1, v_ax2], linear_bias196[v_ax2])
+ T.writes(var_T_add_intermediate[v_ax0, v_ax1, v_ax2])
+ var_T_add_intermediate[v_ax0, v_ax1, v_ax2] = var_matmul_intermediate[v_ax0, v_ax1, v_ax2] + linear_bias196[v_ax2]
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(10240)):
+ with T.block("T_multiply"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_T_add_intermediate[v_ax0, v_ax1, v_ax2])
+ T.writes(T_multiply[v_ax0, v_ax1, v_ax2])
+ T_multiply[v_ax0, v_ax1, v_ax2] = var_T_add_intermediate[v_ax0, v_ax1, v_ax2] * T.float32(0.70710678118654757)
+ for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), T.int64(10240)):
+ with T.block("compute"):
+ v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(T_multiply[v_i0, v_i1, v_i2])
+ T.writes(compute[v_i0, v_i1, v_i2])
+ compute[v_i0, v_i1, v_i2] = T.erf(T_multiply[v_i0, v_i1, v_i2])
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(10240)):
+ with T.block("T_multiply_1"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(compute[v_ax0, v_ax1, v_ax2])
+ T.writes(T_multiply_1[v_ax0, v_ax1, v_ax2])
+ T_multiply_1[v_ax0, v_ax1, v_ax2] = compute[v_ax0, v_ax1, v_ax2] * T.float32(0.5)
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(10240)):
+ with T.block("T_add_1"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(T_multiply_1[v_ax0, v_ax1, v_ax2])
+ T.writes(T_add[v_ax0, v_ax1, v_ax2])
+ T_add[v_ax0, v_ax1, v_ax2] = T.float32(0.5) + T_multiply_1[v_ax0, v_ax1, v_ax2]
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(10240)):
+ with T.block("T_multiply_2"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_T_add_intermediate[v_ax0, v_ax1, v_ax2], T_add[v_ax0, v_ax1, v_ax2])
+ T.writes(var_T_multiply_intermediate[v_ax0, v_ax1, v_ax2])
+ var_T_multiply_intermediate[v_ax0, v_ax1, v_ax2] = var_T_add_intermediate[v_ax0, v_ax1, v_ax2] * T_add[v_ax0, v_ax1, v_ax2]
+ for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), T.int64(10240)):
+ with T.block("compute_1"):
+ v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(var_T_multiply_intermediate[v_i0, v_i1, v_i2])
+ T.writes(p_output0_intermediate[v_i0, v_i1, v_i2])
+ p_output0_intermediate[v_i0, v_i1, v_i2] = T.Cast("float16", var_T_multiply_intermediate[v_i0, v_i1, v_i2])
+
+ @T.prim_func
+ def func5(lv1393: T.Buffer((T.int64(1280), T.int64(2560)), "uint32"), lv1394: T.Buffer((T.int64(320), T.int64(2560)), "float16"), lv2121: T.Buffer((T.int64(1), T.int64(1), T.int64(10240)), "float16"), linear_bias197: T.Buffer((T.int64(2560),), "float32"), lv329: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16"), p_output0_intermediate: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ var_decode_intermediate = T.alloc_buffer((T.int64(10240), T.int64(2560)), "float16")
+ var_matmul_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)))
+ var_T_add_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)))
+ var_compute_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")
+ var_compute_intermediate_1 = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")
+ for i, j in T.grid(T.int64(10240), T.int64(2560)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(lv1393[v_i // T.int64(8), v_j], lv1394[v_i // T.int64(32), v_j])
+ T.writes(var_decode_intermediate[v_i, v_j])
+ var_decode_intermediate[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv1393[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv1394[v_i // T.int64(32), v_j]
+ for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), T.int64(2560), T.int64(10240)):
+ with T.block("matmul"):
+ v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
+ T.reads(lv2121[v_i0, v_i1, v_k], var_decode_intermediate[v_k, v_i2])
+ T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
+ with T.init():
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float32(0)
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + T.Cast("float32", lv2121[v_i0, v_i1, v_k]) * T.Cast("float32", var_decode_intermediate[v_k, v_i2])
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("T_add"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_matmul_intermediate[v_ax0, v_ax1, v_ax2], linear_bias197[v_ax2])
+ T.writes(var_T_add_intermediate[v_ax0, v_ax1, v_ax2])
+ var_T_add_intermediate[v_ax0, v_ax1, v_ax2] = var_matmul_intermediate[v_ax0, v_ax1, v_ax2] + linear_bias197[v_ax2]
+ for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("compute"):
+ v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(var_T_add_intermediate[v_i0, v_i1, v_i2])
+ T.writes(var_compute_intermediate[v_i0, v_i1, v_i2])
+ var_compute_intermediate[v_i0, v_i1, v_i2] = T.Cast("float16", var_T_add_intermediate[v_i0, v_i1, v_i2])
+ for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("compute_1"):
+ v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(var_compute_intermediate[v_i0, v_i1, v_i2])
+ T.writes(var_compute_intermediate_1[v_i0, v_i1, v_i2])
+ var_compute_intermediate_1[v_i0, v_i1, v_i2] = var_compute_intermediate[v_i0, v_i1, v_i2]
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("T_add_1"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_compute_intermediate_1[v_ax0, v_ax1, v_ax2], lv329[v_ax0, v_ax1, v_ax2])
+ T.writes(p_output0_intermediate[v_ax0, v_ax1, v_ax2])
+ p_output0_intermediate[v_ax0, v_ax1, v_ax2] = var_compute_intermediate_1[v_ax0, v_ax1, v_ax2] + lv329[v_ax0, v_ax1, v_ax2]
+
+ @T.prim_func
+ def func6(lv2509: T.Buffer((T.int64(1280), T.int64(2560)), "uint32"), lv2510: T.Buffer((T.int64(320), T.int64(2560)), "float16"), lv4105: T.Buffer((T.int64(1), T.int64(1), T.int64(10240)), "float16"), linear_bias383: T.Buffer((T.int64(2560),), "float32"), lv701: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16"), p_output0_intermediate: T.Buffer((T.int64(1), T.int64(1), T.int64(2560)), "float32")):
+ T.func_attr({"tir.noalias": T.bool(True)})
+ # with T.block("root"):
+ var_decode_intermediate = T.alloc_buffer((T.int64(10240), T.int64(2560)), "float16")
+ var_matmul_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)))
+ var_T_add_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)))
+ var_compute_intermediate = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")
+ var_compute_intermediate_1 = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")
+ var_T_add_intermediate_1 = T.alloc_buffer((T.int64(1), T.int64(1), T.int64(2560)), "float16")
+ for i, j in T.grid(T.int64(10240), T.int64(2560)):
+ with T.block("decode"):
+ v_i, v_j = T.axis.remap("SS", [i, j])
+ T.reads(lv2509[v_i // T.int64(8), v_j], lv2510[v_i // T.int64(32), v_j])
+ T.writes(var_decode_intermediate[v_i, v_j])
+ var_decode_intermediate[v_i, v_j] = (T.Cast("float16", T.bitwise_and(T.shift_right(lv2509[v_i // T.int64(8), v_j], T.Cast("uint32", v_i % T.int64(8)) * T.uint32(4)), T.uint32(15))) - T.float16(7)) * lv2510[v_i // T.int64(32), v_j]
+ for i0, i1, i2, k in T.grid(T.int64(1), T.int64(1), T.int64(2560), T.int64(10240)):
+ with T.block("matmul"):
+ v_i0, v_i1, v_i2, v_k = T.axis.remap("SSSR", [i0, i1, i2, k])
+ T.reads(lv4105[v_i0, v_i1, v_k], var_decode_intermediate[v_k, v_i2])
+ T.writes(var_matmul_intermediate[v_i0, v_i1, v_i2])
+ with T.init():
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = T.float32(0)
+ var_matmul_intermediate[v_i0, v_i1, v_i2] = var_matmul_intermediate[v_i0, v_i1, v_i2] + T.Cast("float32", lv4105[v_i0, v_i1, v_k]) * T.Cast("float32", var_decode_intermediate[v_k, v_i2])
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("T_add"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_matmul_intermediate[v_ax0, v_ax1, v_ax2], linear_bias383[v_ax2])
+ T.writes(var_T_add_intermediate[v_ax0, v_ax1, v_ax2])
+ var_T_add_intermediate[v_ax0, v_ax1, v_ax2] = var_matmul_intermediate[v_ax0, v_ax1, v_ax2] + linear_bias383[v_ax2]
+ for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("compute"):
+ v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(var_T_add_intermediate[v_i0, v_i1, v_i2])
+ T.writes(var_compute_intermediate[v_i0, v_i1, v_i2])
+ var_compute_intermediate[v_i0, v_i1, v_i2] = T.Cast("float16", var_T_add_intermediate[v_i0, v_i1, v_i2])
+ for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("compute_1"):
+ v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(var_compute_intermediate[v_i0, v_i1, v_i2])
+ T.writes(var_compute_intermediate_1[v_i0, v_i1, v_i2])
+ var_compute_intermediate_1[v_i0, v_i1, v_i2] = var_compute_intermediate[v_i0, v_i1, v_i2]
+ for ax0, ax1, ax2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("T_add_1"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(var_compute_intermediate_1[v_ax0, v_ax1, v_ax2], lv701[v_ax0, v_ax1, v_ax2])
+ T.writes(var_T_add_intermediate_1[v_ax0, v_ax1, v_ax2])
+ var_T_add_intermediate_1[v_ax0, v_ax1, v_ax2] = var_compute_intermediate_1[v_ax0, v_ax1, v_ax2] + lv701[v_ax0, v_ax1, v_ax2]
+ for i0, i1, i2 in T.grid(T.int64(1), T.int64(1), T.int64(2560)):
+ with T.block("compute_2"):
+ v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
Review Comment:
Would be good to delete this for now, and as we add new test cases like `test_normalize.py` we add these modules as before/after in each test file.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on a diff in pull request #15141: [WIP] Scaffolding DLight
Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen commented on code in PR #15141:
URL: https://github.com/apache/tvm/pull/15141#discussion_r1238660532
##########
main.py:
##########
@@ -0,0 +1,267 @@
+# pylint: disable=missing-docstring
+from typing import List, Optional
+
+from tvm import IRModule
+from tvm import meta_schedule as ms
+from tvm import tir
+from tvm.dlight import ScheduleGenerator, ScheduleRule, auto_inline_consumers
+
+
+class Decode(ScheduleRule):
+ def __init__(self):
+ ...
+
+ def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
+ pass
+
+ def clone(self) -> ScheduleRule:
+ return Decode()
+
+ def apply(self, mod: IRModule) -> Optional[tir.Schedule]: # pylint: disable=too-many-locals
+ sch = tir.Schedule(mod)
+ try:
+ decode = sch.get_block("decode")
+ except: # pylint: disable=bare-except
+ return None
+ len_tx: int = 8
+ len_ty: int = 8
+ len_yi: int = 1
+ len_yc: int = 8
+
+ # Step 1. Tile the decoding
+ i, j = sch.get_loops(decode)
+ by, ty, yi, yc = sch.split( # pylint: disable=invalid-name
+ i, factors=[None, len_ty, len_yi, len_yc]
+ )
+ bx, tx = sch.split(j, factors=[None, len_tx]) # pylint: disable=invalid-name
+ sch.reorder(by, bx, ty, tx, yi, yc)
+ sch.bind(by, "blockIdx.y")
+ sch.bind(bx, "blockIdx.x")
+ sch.bind(ty, "threadIdx.y")
+ sch.bind(tx, "threadIdx.x")
+ sch.unroll(yc)
+ # Step 2. Cache results in shared memory
+ rb = sch.cache_write(decode, 0, "shared") # pylint: disable=invalid-name
+ consumers = sch.get_consumers(rb)
+ if consumers:
+ (consumer,) = consumers
+ auto_inline_consumers(sch, consumer)
+ sch.compute_inline(rb)
+ rb = consumer # pylint: disable=invalid-name
+ # Step 3. Schedule the shared memory write back
+ sch.reverse_compute_at(rb, bx, preserve_unit_loops=True)
+ loop = sch.fuse(*sch.get_loops(rb)[-2:])
+ _, ty, tx, vec = sch.split( # pylint: disable=invalid-name
+ loop, factors=[None, len_ty, len_tx, 4]
+ )
+ sch.bind(ty, "threadIdx.y")
+ sch.bind(tx, "threadIdx.x")
+ sch.vectorize(vec)
+ sch.storage_align(decode, buffer_index=0, axis=0, factor=32, offset=1)
+ return sch
+
+
+class DecodeGemv(ScheduleRule):
+ def __init__(self):
+ ...
+
+ def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
+ pass
+
+ def clone(self) -> ScheduleRule:
+ return DecodeGemv()
+
+ def apply(self, mod: IRModule) -> Optional[tir.Schedule]: # pylint: disable=too-many-locals
+ sch = tir.Schedule(mod)
+ try:
+ gemv = sch.get_block("matmul")
+ decode = sch.get_block("decode")
+ except: # pylint: disable=bare-except
+ return None
+ len_vx: int = 2
+ len_tx: int = 64
+ len_km: int = 2
+ len_ki: int = 1 * 8
+ # Step 1. Schedule GEMV
+ # [b=1, i=1, j, k]
+ # split j => [b=1, i=1, (bx, vx, tx), k]
+ # fuse (b, i, bx) => [bx, vx, tx, (k)]
+ # split k => [bx, vx, tx, (ko, k_m, ki * 8)]
+ rb = sch.cache_write(gemv, 0, "local") # pylint: disable=invalid-name
+ b, i, j, k = sch.get_loops(gemv) # pylint: disable=invalid-name
+ assert sch.get(b).extent.value == 1
+ assert sch.get(i).extent.value == 1
+ bx, vx, tx = sch.split(j, [None, len_vx, len_tx]) # pylint: disable=invalid-name
+ bx = sch.fuse(b, i, bx) # pylint: disable=invalid-name
+ k_o, k_m, k_i = sch.split(k, [None, len_km, len_ki])
+ sch.bind(bx, thread_axis="blockIdx.x")
+ sch.bind(vx, thread_axis="vthread.x")
+ sch.bind(tx, thread_axis="threadIdx.x")
+ sch.reorder(bx, vx, tx, k_o, k_m, k_i)
+ sch.unroll(k_i)
+ # Step 2. Schedule decode: move to under threadIdx.x and fetch separately for each thread
+ sch.compute_at(decode, k_m, preserve_unit_loops=True)
+ sch.set_scope(decode, 0, "local")
+ _, unroll = sch.split(sch.get_loops(decode)[-2], [None, 8])
+ sch.unroll(unroll)
+
+ # Step 3. Cooperative fetch GEMV
+ def cooperative_fetch(block, tx): # pylint: disable=invalid-name
+ block = sch.cache_read(block, 0, "shared")
+ sch.compute_at(block, tx, preserve_unit_loops=True)
+ loop = sch.fuse(*sch.get_loops(block)[-2:])
+ len_vector = sch.sample_categorical(
+ [1, 2, 3, 4],
+ probs=[0.25, 0.25, 0.25, 0.25],
+ )
+ _, tx, vec = sch.split(loop, [None, len_tx, len_vector])
+ sch.bind(tx, thread_axis="threadIdx.x")
+ sch.vectorize(vec)
+ sch.storage_align(block, buffer_index=0, axis=-2, factor=32, offset=8)
+
+ cooperative_fetch(gemv, k_o)
+ # Step 4. Schedule epilogue
+ auto_inline_consumers(sch, rb)
+ sch.reverse_compute_at(rb, tx, preserve_unit_loops=True)
+ # Step 5. Postprocess: decompose reduction
+ sch.decompose_reduction(gemv, k_o)
+ return [sch]
+
+
+class Normalization(ScheduleRule):
+ def __init__(self):
+ ...
+
+ def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
+ pass
+
+ def clone(self) -> ScheduleRule:
+ return Normalization()
+
+ def apply(self, mod: IRModule) -> Optional[tir.Schedule]: # pylint: disable=too-many-locals
+ sch = tir.Schedule(mod)
+ b_reduce = None
+ for name in ["Ared_temp", "A_red_temp"]:
+ try:
+ b_reduce = sch.get_block(name)
+ except:
+ continue
+ else:
+ break
+ if b_reduce is None:
+ return None
+ len_tx: int = 256
+ unroll_depth: int = 256
+
+ (b_spatial,) = sch.get_consumers(b_reduce)
+ loops = sch.get_loops(b_spatial)
+ bx = sch.fuse(*loops[:-1]) # pylint: disable=invalid-name
+ _, tx = sch.split(loops[-1], [None, len_tx]) # pylint: disable=invalid-name
+ sch.bind(bx, "blockIdx.x")
+ sch.bind(tx, "threadIdx.x")
+
+ for i, _ in enumerate(sch.get(b_reduce).writes):
+ sch.set_scope(b_reduce, buffer_index=i, storage_scope="shared")
+ sch.compute_at(b_reduce, bx, preserve_unit_loops=True)
+ _, tx = sch.split( # pylint: disable=invalid-name
+ sch.get_loops(b_reduce)[-1],
+ [None, len_tx],
+ )
+ sch.bind(tx, "threadIdx.x")
+ auto_inline_consumers(sch, b_spatial)
+ sch.annotate(bx, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
+ sch.annotate(bx, ann_key="pragma_unroll_explicit", ann_val=1)
+
+ return sch
+
+
+class Softmax(ScheduleRule):
+ def __init__(self):
+ ...
+
+ def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
+ pass
+
+ def clone(self) -> ScheduleRule:
+ return Softmax()
+
+ def apply(self, mod: IRModule) -> Optional[tir.Schedule]: # pylint: disable=too-many-locals
+ sch = tir.Schedule(mod)
+ try:
+ b_reduce_0 = sch.get_block("T_softmax_maxelem")
+ b_reduce_1 = sch.get_block("T_softmax_expsum")
+ b_spatial = sch.get_block("T_softmax_norm")
+ except: # pylint: disable=bare-except
+ return None
+
+ len_tx: int = 256
+ unroll_depth: int = 256
+
+ sch.compute_inline(sch.get_producers(b_reduce_1)[0])
+
+ loops = sch.get_loops(b_spatial)
+ bx = sch.fuse(*loops[:-1]) # pylint: disable=invalid-name
+ _, tx = sch.split(loops[-1], [None, len_tx]) # pylint: disable=invalid-name
+ sch.bind(bx, "blockIdx.x")
+ sch.bind(tx, "threadIdx.x")
+
+ sch.set_scope(b_reduce_1, buffer_index=0, storage_scope="shared")
+ sch.compute_at(b_reduce_1, bx, preserve_unit_loops=True)
+ _, tx = sch.split( # pylint: disable=invalid-name
+ sch.get_loops(b_reduce_1)[-1],
+ [None, len_tx],
+ )
+ sch.bind(tx, "threadIdx.x")
+
+ sch.set_scope(b_reduce_0, buffer_index=0, storage_scope="shared")
+ sch.compute_at(b_reduce_0, bx, preserve_unit_loops=True)
+ _, tx = sch.split( # pylint: disable=invalid-name
+ sch.get_loops(b_reduce_0)[-1],
+ [None, len_tx],
+ )
+ sch.bind(tx, "threadIdx.x")
+
+ sch.annotate(bx, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
+ sch.annotate(bx, ann_key="pragma_unroll_explicit", ann_val=1)
+
+ auto_inline_consumers(sch, b_spatial)
+ return [sch]
+
+
+def main():
+ from tvm.dlight.testing import mod_decode, mod_decode_gemv, mod_norm, mod_softmax
+
+ gen = ScheduleGenerator(
+ rules=[
+ DecodeGemv(),
+ Decode(),
+ Normalization(),
+ Softmax(),
+ ]
+ )
+
+ for py_mod in [
+ mod_decode,
+ mod_decode_gemv,
+ mod_norm,
+ mod_softmax, # Needs to upstream `compute-inline`
+ ]:
+ i = 1
+ while True:
+ try:
+ func = py_mod.Module[f"func{i}"]
+ except: # pylint: disable=bare-except
+ break
+ else:
+ print(f"Working on {py_mod}::func{i}")
+ i += 1
+ mod = IRModule.from_expr(func.with_attr("global_symbol", "main"))
+ schedules = gen.generate_design_space(mod)
Review Comment:
Let us structure the testcase UT same as Before/After test, e.g. one softmax function before and after with structural equality test. (gen and apply)
Likely most of the modules in dlight testing can move to `tests/dlight/` (that can be added to CI)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on a diff in pull request #15141: [WIP] Scaffolding DLight
Posted by "tqchen (via GitHub)" <gi...@apache.org>.
tqchen commented on code in PR #15141:
URL: https://github.com/apache/tvm/pull/15141#discussion_r1238660532
##########
main.py:
##########
@@ -0,0 +1,267 @@
+# pylint: disable=missing-docstring
+from typing import List, Optional
+
+from tvm import IRModule
+from tvm import meta_schedule as ms
+from tvm import tir
+from tvm.dlight import ScheduleGenerator, ScheduleRule, auto_inline_consumers
+
+
+class Decode(ScheduleRule):
+ def __init__(self):
+ ...
+
+ def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
+ pass
+
+ def clone(self) -> ScheduleRule:
+ return Decode()
+
+ def apply(self, mod: IRModule) -> Optional[tir.Schedule]: # pylint: disable=too-many-locals
+ sch = tir.Schedule(mod)
+ try:
+ decode = sch.get_block("decode")
+ except: # pylint: disable=bare-except
+ return None
+ len_tx: int = 8
+ len_ty: int = 8
+ len_yi: int = 1
+ len_yc: int = 8
+
+ # Step 1. Tile the decoding
+ i, j = sch.get_loops(decode)
+ by, ty, yi, yc = sch.split( # pylint: disable=invalid-name
+ i, factors=[None, len_ty, len_yi, len_yc]
+ )
+ bx, tx = sch.split(j, factors=[None, len_tx]) # pylint: disable=invalid-name
+ sch.reorder(by, bx, ty, tx, yi, yc)
+ sch.bind(by, "blockIdx.y")
+ sch.bind(bx, "blockIdx.x")
+ sch.bind(ty, "threadIdx.y")
+ sch.bind(tx, "threadIdx.x")
+ sch.unroll(yc)
+ # Step 2. Cache results in shared memory
+ rb = sch.cache_write(decode, 0, "shared") # pylint: disable=invalid-name
+ consumers = sch.get_consumers(rb)
+ if consumers:
+ (consumer,) = consumers
+ auto_inline_consumers(sch, consumer)
+ sch.compute_inline(rb)
+ rb = consumer # pylint: disable=invalid-name
+ # Step 3. Schedule the shared memory write back
+ sch.reverse_compute_at(rb, bx, preserve_unit_loops=True)
+ loop = sch.fuse(*sch.get_loops(rb)[-2:])
+ _, ty, tx, vec = sch.split( # pylint: disable=invalid-name
+ loop, factors=[None, len_ty, len_tx, 4]
+ )
+ sch.bind(ty, "threadIdx.y")
+ sch.bind(tx, "threadIdx.x")
+ sch.vectorize(vec)
+ sch.storage_align(decode, buffer_index=0, axis=0, factor=32, offset=1)
+ return sch
+
+
+class DecodeGemv(ScheduleRule):
+ def __init__(self):
+ ...
+
+ def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
+ pass
+
+ def clone(self) -> ScheduleRule:
+ return DecodeGemv()
+
+ def apply(self, mod: IRModule) -> Optional[tir.Schedule]: # pylint: disable=too-many-locals
+ sch = tir.Schedule(mod)
+ try:
+ gemv = sch.get_block("matmul")
+ decode = sch.get_block("decode")
+ except: # pylint: disable=bare-except
+ return None
+ len_vx: int = 2
+ len_tx: int = 64
+ len_km: int = 2
+ len_ki: int = 1 * 8
+ # Step 1. Schedule GEMV
+ # [b=1, i=1, j, k]
+ # split j => [b=1, i=1, (bx, vx, tx), k]
+ # fuse (b, i, bx) => [bx, vx, tx, (k)]
+ # split k => [bx, vx, tx, (ko, k_m, ki * 8)]
+ rb = sch.cache_write(gemv, 0, "local") # pylint: disable=invalid-name
+ b, i, j, k = sch.get_loops(gemv) # pylint: disable=invalid-name
+ assert sch.get(b).extent.value == 1
+ assert sch.get(i).extent.value == 1
+ bx, vx, tx = sch.split(j, [None, len_vx, len_tx]) # pylint: disable=invalid-name
+ bx = sch.fuse(b, i, bx) # pylint: disable=invalid-name
+ k_o, k_m, k_i = sch.split(k, [None, len_km, len_ki])
+ sch.bind(bx, thread_axis="blockIdx.x")
+ sch.bind(vx, thread_axis="vthread.x")
+ sch.bind(tx, thread_axis="threadIdx.x")
+ sch.reorder(bx, vx, tx, k_o, k_m, k_i)
+ sch.unroll(k_i)
+ # Step 2. Schedule decode: move to under threadIdx.x and fetch separately for each thread
+ sch.compute_at(decode, k_m, preserve_unit_loops=True)
+ sch.set_scope(decode, 0, "local")
+ _, unroll = sch.split(sch.get_loops(decode)[-2], [None, 8])
+ sch.unroll(unroll)
+
+ # Step 3. Cooperative fetch GEMV
+ def cooperative_fetch(block, tx): # pylint: disable=invalid-name
+ block = sch.cache_read(block, 0, "shared")
+ sch.compute_at(block, tx, preserve_unit_loops=True)
+ loop = sch.fuse(*sch.get_loops(block)[-2:])
+ len_vector = sch.sample_categorical(
+ [1, 2, 3, 4],
+ probs=[0.25, 0.25, 0.25, 0.25],
+ )
+ _, tx, vec = sch.split(loop, [None, len_tx, len_vector])
+ sch.bind(tx, thread_axis="threadIdx.x")
+ sch.vectorize(vec)
+ sch.storage_align(block, buffer_index=0, axis=-2, factor=32, offset=8)
+
+ cooperative_fetch(gemv, k_o)
+ # Step 4. Schedule epilogue
+ auto_inline_consumers(sch, rb)
+ sch.reverse_compute_at(rb, tx, preserve_unit_loops=True)
+ # Step 5. Postprocess: decompose reduction
+ sch.decompose_reduction(gemv, k_o)
+ return [sch]
+
+
+class Normalization(ScheduleRule):
+ def __init__(self):
+ ...
+
+ def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
+ pass
+
+ def clone(self) -> ScheduleRule:
+ return Normalization()
+
+ def apply(self, mod: IRModule) -> Optional[tir.Schedule]: # pylint: disable=too-many-locals
+ sch = tir.Schedule(mod)
+ b_reduce = None
+ for name in ["Ared_temp", "A_red_temp"]:
+ try:
+ b_reduce = sch.get_block(name)
+ except:
+ continue
+ else:
+ break
+ if b_reduce is None:
+ return None
+ len_tx: int = 256
+ unroll_depth: int = 256
+
+ (b_spatial,) = sch.get_consumers(b_reduce)
+ loops = sch.get_loops(b_spatial)
+ bx = sch.fuse(*loops[:-1]) # pylint: disable=invalid-name
+ _, tx = sch.split(loops[-1], [None, len_tx]) # pylint: disable=invalid-name
+ sch.bind(bx, "blockIdx.x")
+ sch.bind(tx, "threadIdx.x")
+
+ for i, _ in enumerate(sch.get(b_reduce).writes):
+ sch.set_scope(b_reduce, buffer_index=i, storage_scope="shared")
+ sch.compute_at(b_reduce, bx, preserve_unit_loops=True)
+ _, tx = sch.split( # pylint: disable=invalid-name
+ sch.get_loops(b_reduce)[-1],
+ [None, len_tx],
+ )
+ sch.bind(tx, "threadIdx.x")
+ auto_inline_consumers(sch, b_spatial)
+ sch.annotate(bx, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
+ sch.annotate(bx, ann_key="pragma_unroll_explicit", ann_val=1)
+
+ return sch
+
+
+class Softmax(ScheduleRule):
+ def __init__(self):
+ ...
+
+ def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
+ pass
+
+ def clone(self) -> ScheduleRule:
+ return Softmax()
+
+ def apply(self, mod: IRModule) -> Optional[tir.Schedule]: # pylint: disable=too-many-locals
+ sch = tir.Schedule(mod)
+ try:
+ b_reduce_0 = sch.get_block("T_softmax_maxelem")
+ b_reduce_1 = sch.get_block("T_softmax_expsum")
+ b_spatial = sch.get_block("T_softmax_norm")
+ except: # pylint: disable=bare-except
+ return None
+
+ len_tx: int = 256
+ unroll_depth: int = 256
+
+ sch.compute_inline(sch.get_producers(b_reduce_1)[0])
+
+ loops = sch.get_loops(b_spatial)
+ bx = sch.fuse(*loops[:-1]) # pylint: disable=invalid-name
+ _, tx = sch.split(loops[-1], [None, len_tx]) # pylint: disable=invalid-name
+ sch.bind(bx, "blockIdx.x")
+ sch.bind(tx, "threadIdx.x")
+
+ sch.set_scope(b_reduce_1, buffer_index=0, storage_scope="shared")
+ sch.compute_at(b_reduce_1, bx, preserve_unit_loops=True)
+ _, tx = sch.split( # pylint: disable=invalid-name
+ sch.get_loops(b_reduce_1)[-1],
+ [None, len_tx],
+ )
+ sch.bind(tx, "threadIdx.x")
+
+ sch.set_scope(b_reduce_0, buffer_index=0, storage_scope="shared")
+ sch.compute_at(b_reduce_0, bx, preserve_unit_loops=True)
+ _, tx = sch.split( # pylint: disable=invalid-name
+ sch.get_loops(b_reduce_0)[-1],
+ [None, len_tx],
+ )
+ sch.bind(tx, "threadIdx.x")
+
+ sch.annotate(bx, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
+ sch.annotate(bx, ann_key="pragma_unroll_explicit", ann_val=1)
+
+ auto_inline_consumers(sch, b_spatial)
+ return [sch]
+
+
+def main():
+ from tvm.dlight.testing import mod_decode, mod_decode_gemv, mod_norm, mod_softmax
+
+ gen = ScheduleGenerator(
+ rules=[
+ DecodeGemv(),
+ Decode(),
+ Normalization(),
+ Softmax(),
+ ]
+ )
+
+ for py_mod in [
+ mod_decode,
+ mod_decode_gemv,
+ mod_norm,
+ mod_softmax, # Needs to upstream `compute-inline`
+ ]:
+ i = 1
+ while True:
+ try:
+ func = py_mod.Module[f"func{i}"]
+ except: # pylint: disable=bare-except
+ break
+ else:
+ print(f"Working on {py_mod}::func{i}")
+ i += 1
+ mod = IRModule.from_expr(func.with_attr("global_symbol", "main"))
+ schedules = gen.generate_design_space(mod)
Review Comment:
Let us structure the testcase UT same as Before/After test, e.g. one softmax function before and after with structural equality test.
Atm we only test that the schedule exists
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao commented on a diff in pull request #15141: [WIP] Scaffolding DLight
Posted by "junrushao (via GitHub)" <gi...@apache.org>.
junrushao commented on code in PR #15141:
URL: https://github.com/apache/tvm/pull/15141#discussion_r1239922425
##########
python/tvm/dlight/testing/modules.py:
##########
@@ -0,0 +1,2621 @@
+from tvm.script import ir as I
+from tvm.script import tir as T
+
+# fmt: off
+
+@I.ir_module
+class Decode:
+
Review Comment:
moved to `tests/python/dlight/test_modules.py`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org