JackWeiw commented on code in PR #15986:
URL: https://github.com/apache/tvm/pull/15986#discussion_r1374209716
##########
tests/python/unittest/test_tir_transform_lower_opaque_block.py:
##########
@@ -250,6 +250,34 @@ def transformed_strided_buffer_func(
C[i0 * 4 + i1, j] = B[i1, j] * T.float32(2)
[email protected]_func
+def compacted_symbolic_strided_buffer_func(a: T.handle) -> None:
+ n = T.int64()
+ A = T.match_buffer(a, (1, n, 10240), "float32")
+ for i, j, k in T.grid(((n + 63) // 64 * 4 + 7) // 8, 2, 160):
+ with T.block(""):
+ T.reads(A[0, i * 128 + j * 32:i * 128 + j * 32 + 96, k * 64:k * 64
+ 64])
+ A_pad_shared_dyn = T.alloc_buffer((1, T.min((n + 63) // 64 * 64,
96), 64), "float32", strides=(72 * T.min((n + 63) // 64 * 64, 96), 72, 1),
scope="shared.dyn")
+ for ax0, ax1 in T.grid(96, 64):
+ with T.block("A_pad_shared.dyn"):
+ T.where(i * 128 + j * 32 + ax0 < (n + 63) // 64 * 64)
+ T.reads(A[0, i * 128 + j * 32 + ax0, k * 64 + ax1])
+ T.writes(A_pad_shared_dyn[0, ax0, ax1])
+ A_pad_shared_dyn[0, ax0, ax1] = T.if_then_else(i * 128 + j
* 32 + ax0 < n, A[0, i * 128 + j * 32 + ax0, k * 64 + ax1], T.float16(0))
+
+
[email protected]_func
+def transformed_symbolic_strided_buffer_func(a: T.handle):
+ n = T.int64()
+ A = T.match_buffer(a, (1, n, 10240))
+ for i, j, k in T.grid(((n + T.int64(63)) // T.int64(64) * T.int64(4) +
T.int64(7)) // T.int64(8), 2, 160):
Review Comment:
Thank u for the adivice, i've modified it using` T.int32` and pulled out
into `padded_size = T.meta_var(T.min((n + T.int64(63)) // T.int64(64) *
T.int64(64))` in the test case.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
