You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2021/03/08 21:00:50 UTC
[GitHub] [tvm] mbrookhart opened a new pull request #7611: Mergepath sort with odd-even block sort
mbrookhart opened a new pull request #7611:
URL: https://github.com/apache/tvm/pull/7611
Building on @masahi's work to add a while loop to TIR, this PR implements the [MergePath](https://arxiv.org/abs/1406.2628) algorithm to parallelize the later stages of mergesort. It also implements a stable in-shared-memory quicksort to do small block sorting before mergesort for better speed. All told, this optimization gives us dramatic speedups on both AMD and Nvidia GPUs, see below. cc @Laurawly @zhiics @icemelon9 @csullivan @tkonolige
I had to add a path for skipping thread planning to get handed-threaded odd-even-tranpose sort to work, cc @tqchen @junrushao1994
```
thrust_R9-Nano mergepath_R9-Nano thrust_1070Ti mergepath_1070Ti thrust_V100 mergepath_V100
[1000, 1, 1] 0.60 ms 0.15 ms 0.05 ms 0.05 ms 0.07 ms 0.04 ms
[1, 1000, 1] 0.59 ms 0.14 ms 0.05 ms 0.05 ms 0.07 ms 0.04 ms
[1, 1, 1000] 0.71 ms 0.13 ms 0.04 ms 0.05 ms 0.06 ms 0.03 ms
[1000, 10, 10] 1.65 ms 0.25 ms 0.96 ms 0.27 ms 1.81 ms 0.08 ms
[10, 1000, 10] 1.65 ms 0.25 ms 0.80 ms 0.26 ms 1.64 ms 0.08 ms
[10, 10, 1000] 1.63 ms 0.23 ms 0.95 ms 0.25 ms 1.82 ms 0.07 ms
[1000, 100, 100] 29.29 ms 18.59 ms 47.15 ms 33.82 ms 18.74 ms 7.43 ms
[100, 1000, 100] 28.63 ms 16.35 ms 37.36 ms 23.57 ms 17.44 ms 5.28 ms
[100, 100, 1000] 29.62 ms 15.13 ms 36.33 ms 21.11 ms 18.04 ms 4.64 ms
[10000, 1, 1] 0.78 ms 0.25 ms 0.14 ms 0.10 ms 0.46 ms 0.08 ms
[1, 10000, 1] 0.79 ms 0.25 ms 0.14 ms 0.10 ms 0.45 ms 0.08 ms
[1, 1, 10000] 0.76 ms 0.24 ms 0.13 ms 0.09 ms 0.45 ms 0.08 ms
[10000, 10, 10] 5.27 ms 2.68 ms 4.58 ms 3.36 ms 3.78 ms 0.78 ms
[10, 10000, 10] 5.03 ms 2.56 ms 4.47 ms 3.33 ms 3.78 ms 0.73 ms
[10, 10, 10000] 5.34 ms 2.48 ms 4.34 ms 3.08 ms 3.76 ms 0.70 ms
[10000, 100, 100] Out of memory 338.65 ms 467.08 ms 437.17 ms 145.14 ms 96.20 ms
[100, 10000, 100] Out of memory 284.24 ms 410.39 ms 380.18 ms 134.70 ms 84.63 ms
[100, 100, 10000] Out of memory 250.24 ms 367.42 ms 305.95 ms 111.16 ms 59.69 ms
[100000, 1, 1] 0.73 ms 0.51 ms 0.16 ms 0.40 ms 0.50 ms 0.14 ms
[1, 100000, 1] 0.73 ms 0.51 ms 0.27 ms 0.40 ms 0.66 ms 0.14 ms
[1, 1, 100000] 0.71 ms 0.50 ms 0.23 ms 0.39 ms 0.69 ms 0.13 ms
[100000, 10, 10] 26.64 ms 33.12 ms 39.23 ms 39.64 ms 18.65 ms 8.48 ms
[10, 100000, 10] 26.80 ms 30.91 ms 37.43 ms 37.88 ms 17.74 ms 7.67 ms
[10, 10, 100000] 25.46 ms 29.75 ms 36.33 ms 35.59 ms 16.49 ms 7.17 ms
[1000000, 1, 1] 1.91 ms 3.68 ms 1.18 ms 4.36 ms 1.00 ms 1.01 ms
[1, 1000000, 1] 1.91 ms 3.68 ms 1.23 ms 4.32 ms 1.00 ms 1.01 ms
[1, 1, 1000000] 1.91 ms 3.64 ms 1.15 ms 4.34 ms 1.01 ms 0.99 ms
[1000000, 10, 10] Out of memory 449.20 ms 410.63 ms 499.69 ms 136.20 ms 115.64 ms
[10, 1000000, 10] Out of memory 385.10 ms 366.33 ms 443.03 ms 116.31 ms 94.27 ms
[10, 10, 1000000] Out of memory 370.27 ms 364.14 ms 419.95 ms 111.16 ms 87.82 ms
[(4507,), 0] 0.70 ms 0.21 ms 1.98 ms 4.32 ms 0.46 ms 0.06 ms
[(1, 122640), 1] 0.78 ms 0.53 ms 0.17 ms 0.45 ms 0.70 ms 0.14 ms
[(1, 120000), 1] 0.73 ms 0.52 ms 0.28 ms 0.44 ms 0.69 ms 0.14 ms
[(1, 30000), 1] 0.74 ms 0.28 ms 0.16 ms 0.14 ms 0.46 ms 0.09 ms
[(1, 7500), 1] 0.73 ms 0.21 ms 0.13 ms 0.08 ms 0.46 ms 0.07 ms
[(1, 1000), 1] 0.65 ms 0.13 ms 0.14 ms 0.07 ms 0.07 ms 0.04 ms
```
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-796019880
On VK, it fails on (122640, 1) workload. rocm works fine.
```
Traceback (most recent call last):
File "vk_test.py", line 146, in <module>
test_argsort()
File "vk_test.py", line 94, in test_argsort
verify_argsort((1, 122640), axis=1, is_ascend=False, dtype=dtype, is_dyn=is_dyn)
File "vk_test.py", line 86, in verify_argsort
tvm.testing.assert_allclose(op_res.asnumpy(), ref_res.astype(dtype), rtol=1e-5)
File "/home/masa/projects/dev/tvm/python/tvm/testing.py", line 82, in assert_allclose
np.testing.assert_allclose(actual, desired, rtol=rtol, atol=atol, verbose=True)
File "/home/masa/anaconda3/lib/python3.8/site-packages/numpy/testing/_private/utils.py", line 1532, in assert_allclose
assert_array_compare(compare, actual, desired, err_msg=str(err_msg),
File "/home/masa/anaconda3/lib/python3.8/site-packages/numpy/testing/_private/utils.py", line 846, in assert_array_compare
raise AssertionError(msg)
AssertionError:
Not equal to tolerance rtol=1e-05, atol=1e-07
Mismatched elements: 122633 / 122640 (100%)
Max absolute difference: 122388
Max relative difference: 84808.
x: array([[116254, 14451, 111687, ..., 64552, 122405, 122011]], dtype=int32)
y: array([[116254, 14451, 111687, ..., 14845, 100388, 53504]], dtype=int32)
```
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] mbrookhart commented on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-794490628
I've got a dynamic topk failure, but I can't reproduce it locally, that test passes on my version of dependencies. Maybe I can run the CI docker locally
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi merged pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi merged pull request #7611:
URL: https://github.com/apache/tvm/pull/7611
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi edited a comment on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-793081245
@mbrookhart This doesn't work on VK/SPIR-V because SPIR-V requires a thread block size to be a compile time constant, see https://github.com/apache/tvm/blob/d7f57532746680732e58ab028d8c3129b9140d3d/src/target/spirv/codegen_spirv.cc#L106
In particular, the following size selection has an issue, because `width` is not a constant:
https://github.com/apache/tvm/blob/d2940ca8e25aa400825b79ac7d03da0ff4543db2/python/tvm/topi/cuda/sort.py#L335-L337
So to workaround this problem, can we either
* Use old TIR sort for VK
* (Better) Use a constant thread block size for VK
?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-796220266
interesting, is it possible to always force single level?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r590689860
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -136,93 +236,223 @@ def compare(a, b):
out = b <= a
return out
- def bottom_up_merge(source, dest, source_idx, dest_idx, start, middle, end, even):
- """
- Merge the two sections of the array assigned to this thread
- """
- # pylint: disable=arguments-out-of-order
- # initialize iterators
- i[0] = start
- j[0] = middle
- # set up indexes
- base_idx = by * size * axis_mul_after + bz
- # iterate over the output loop
- with ib.for_range(0, end - start) as k:
- i_idx = base_idx + i[0] * axis_mul_after
- j_idx = base_idx + j[0] * axis_mul_after
- k_idx = base_idx + (k + start) * axis_mul_after
-
- def swap_values(source, dest, source_idx, dest_idx):
- def assign_i():
- """assign i value to current output"""
- dest[k_idx] = source[i_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[i_idx]
- i[0] += 1
-
- def assign_j():
- """assign j value to current output"""
- dest[k_idx] = source[j_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[j_idx]
- j[0] += 1
-
- ## if both of the iterators are in range
- with ib.if_scope(tvm.tir.all(i[0] < middle, j[0] < end)):
- # compare them and insert whichever is next into the output
- with ib.if_scope(compare(source[i_idx], source[j_idx])):
- assign_i()
- with ib.else_scope():
- assign_j()
- # otherwise, simply copy the remainder of the valid iterator to the output
- with ib.else_scope():
- with ib.if_scope(i[0] < middle):
- assign_i()
- with ib.else_scope():
- assign_j()
+ # Sort the lower levels of the merge using odd-even sort, it's fast for small inputs
+ lower_lim = tvm.tir.generic.cast(
+ tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(block_size, "float64"))), "int64"
+ )
- # Switch which input is the source and which is the destination each iteration
- with ib.if_scope(even):
- swap_values(source, dest, source_idx, dest_idx)
- with ib.else_scope():
- swap_values(dest, source, dest_idx, source_idx)
-
- def mergesort(source, dest, source_idx, dest_idx, size, width, even):
- # calculate the start, mid, and end points of this section
- start[0] = width * tid
- with ib.if_scope(start[0] < size):
- middle[0] = tvm.te.min(start[0] + tvm.tir.indexdiv(width, 2), size)
- end[0] = tvm.te.min(start[0] + width, size)
- ## merge the start->middle and middle->end arrays
- bottom_up_merge(source, dest, source_idx, dest_idx, start[0], middle[0], end[0], even)
-
- lim = tvm.tir.generic.cast(
+ _odd_even_sort(
+ ib,
+ size,
+ axis_mul_before * axis_mul_after,
+ 1,
+ is_ascend,
+ keys,
+ keys_swap,
+ values,
+ values_swap,
+ )
+
+ upper_lim = tvm.tir.generic.cast(
tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(size, "float64"))), "int64"
)
- with ib.for_range(0, lim, dtype="int64") as l2_width:
- width = 2 << l2_width
+
+ def get_merge_begin(source, base_idx, aCount, bCount, aStart, bStart, diag, step_count):
+ first = ib.allocate("int64", (1,), name="first", scope="local")
+ mid = ib.allocate("int64", (1,), name="mid", scope="local")
+ last = ib.allocate("int64", (1,), name="last", scope="local")
+ first[0] = tvm.te.max(0, diag - bCount)
+ last[0] = tvm.te.min(diag, aCount)
+ with ib.while_loop(first[0] < last[0]):
+ mid[0] = (first[0] + last[0]) >> 1
+ a = source[base_idx + (aStart + mid[0])]
+ b = source[base_idx + (bStart + diag - 1 - mid[0])]
+ with ib.if_scope(compare(a, b)):
+ first[0] = mid[0] + 1
+ with ib.else_scope():
+ last[0] = mid[0]
+ return first, last
+
+ def serial_merge(
+ source,
+ dest,
+ source_idx,
+ dest_idx,
+ base_idx,
+ aCount,
+ bCount,
+ aStart,
+ bStart,
+ kStart,
+ diag,
+ step_count,
+ first,
+ last,
+ ):
+ i = ib.allocate("int64", (1,), name="i", scope="local")
+ j = ib.allocate("int64", (1,), name="j", scope="local")
+ i[0] = aStart + first[0]
+ j[0] = bStart + diag - last[0]
+ with ib.for_range(0, tvm.te.min(aCount + bCount - diag, step_count)) as count:
+ i_idx = base_idx + i[0]
+ j_idx = base_idx + j[0]
+ k_idx = base_idx + (kStart + diag + count)
+
+ def assign_i():
+ """assign i value to current output"""
+ dest[k_idx] = source[i_idx]
+ if values is not None:
+ dest_idx[k_idx] = source_idx[i_idx]
+ i[0] += 1
+
+ def assign_j():
+ """assign j value to current output"""
+ dest[k_idx] = source[j_idx]
+ if values is not None:
+ dest_idx[k_idx] = source_idx[j_idx]
+ j[0] += 1
+
+ ## if both of the iterators are in range
+ with ib.if_scope(tvm.tir.all(i[0] < aStart + aCount, j[0] < bStart + bCount)):
+ # compare them and insert whichever is next into the output
+ with ib.if_scope(compare(source[i_idx], source[j_idx])):
+ assign_i()
+ with ib.else_scope():
+ assign_j()
+ # otherwise, simply copy the remainder of the valid iterator to the output
+ with ib.else_scope():
+ with ib.if_scope(i[0] < aStart + aCount):
+ assign_i()
+ with ib.else_scope():
+ assign_j()
+
+ with ib.for_range(0, upper_lim - lower_lim, dtype="int64") as l2_width:
+ width = 2 << (l2_width + lower_lim)
# Define and launch the cuda kernel
with ib.new_scope():
- i = ib.allocate("int64", (1,), name="i", scope="local")
- j = ib.allocate("int64", (1,), name="j", scope="local")
- start = ib.allocate("int64", (1,), name="start", scope="local")
- middle = ib.allocate("int64", (1,), name="middle", scope="local")
- end = ib.allocate("int64", (1,), name="end", scope="local")
- tx = te.thread_axis("threadIdx.x")
- bx = te.thread_axis("blockIdx.x")
- ib.scope_attr(tx, "thread_extent", nthread_tx)
- # Reduce the number of blocks as the work per thread grows
- ib.scope_attr(
- bx,
- "thread_extent",
- tvm.tir.generic.cast(ceil_div(size, width * max_threads), "int32"),
- )
- tid = bx * nthread_tx + tx
Review comment:
You should be able to query the current target from anywhere. You are already doing this when you query the max thread size.
See for example https://github.com/apache/tvm/blob/cf36aa65de80efd84a906058d66b72fb4c23ab16/python/tvm/topi/cuda/nms.py#L632
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi edited a comment on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-796070741
@mbrookhart Also on the failure from `test_any.py`: The last test case tests on 0 size input, and it fails both rocm and vulkan. On VK I get segfault, on rocm I get `Floating point exception (core dumped)`. Maybe you want to check if your code is safe for empty input?
On the old sort IR, rocm works fine. VK still gives segfault, that's another story :)
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi edited a comment on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-793081245
@mbrookhart This doesn't work on VK/SPIR-V because SPIR-V requires a thread block size to be a compile time constant, see https://github.com/apache/tvm/blob/d7f57532746680732e58ab028d8c3129b9140d3d/src/target/spirv/codegen_spirv.cc#L106
In particular, the following size selection has an issue, because `width` is not a constant:
https://github.com/apache/tvm/blob/d2940ca8e25aa400825b79ac7d03da0ff4543db2/python/tvm/topi/cuda/sort.py#L335-L337
So to workaround this problem, can we either
* Use old TIR sort for VK
* (Better) Use a constant thread block size for VK
?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] mbrookhart commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r590563818
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -136,93 +236,223 @@ def compare(a, b):
out = b <= a
return out
- def bottom_up_merge(source, dest, source_idx, dest_idx, start, middle, end, even):
- """
- Merge the two sections of the array assigned to this thread
- """
- # pylint: disable=arguments-out-of-order
- # initialize iterators
- i[0] = start
- j[0] = middle
- # set up indexes
- base_idx = by * size * axis_mul_after + bz
- # iterate over the output loop
- with ib.for_range(0, end - start) as k:
- i_idx = base_idx + i[0] * axis_mul_after
- j_idx = base_idx + j[0] * axis_mul_after
- k_idx = base_idx + (k + start) * axis_mul_after
-
- def swap_values(source, dest, source_idx, dest_idx):
- def assign_i():
- """assign i value to current output"""
- dest[k_idx] = source[i_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[i_idx]
- i[0] += 1
-
- def assign_j():
- """assign j value to current output"""
- dest[k_idx] = source[j_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[j_idx]
- j[0] += 1
-
- ## if both of the iterators are in range
- with ib.if_scope(tvm.tir.all(i[0] < middle, j[0] < end)):
- # compare them and insert whichever is next into the output
- with ib.if_scope(compare(source[i_idx], source[j_idx])):
- assign_i()
- with ib.else_scope():
- assign_j()
- # otherwise, simply copy the remainder of the valid iterator to the output
- with ib.else_scope():
- with ib.if_scope(i[0] < middle):
- assign_i()
- with ib.else_scope():
- assign_j()
+ # Sort the lower levels of the merge using odd-even sort, it's fast for small inputs
+ lower_lim = tvm.tir.generic.cast(
+ tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(block_size, "float64"))), "int64"
+ )
- # Switch which input is the source and which is the destination each iteration
- with ib.if_scope(even):
- swap_values(source, dest, source_idx, dest_idx)
- with ib.else_scope():
- swap_values(dest, source, dest_idx, source_idx)
-
- def mergesort(source, dest, source_idx, dest_idx, size, width, even):
- # calculate the start, mid, and end points of this section
- start[0] = width * tid
- with ib.if_scope(start[0] < size):
- middle[0] = tvm.te.min(start[0] + tvm.tir.indexdiv(width, 2), size)
- end[0] = tvm.te.min(start[0] + width, size)
- ## merge the start->middle and middle->end arrays
- bottom_up_merge(source, dest, source_idx, dest_idx, start[0], middle[0], end[0], even)
-
- lim = tvm.tir.generic.cast(
+ _odd_even_sort(
+ ib,
+ size,
+ axis_mul_before * axis_mul_after,
+ 1,
+ is_ascend,
+ keys,
+ keys_swap,
+ values,
+ values_swap,
+ )
+
+ upper_lim = tvm.tir.generic.cast(
tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(size, "float64"))), "int64"
)
- with ib.for_range(0, lim, dtype="int64") as l2_width:
- width = 2 << l2_width
+
+ def get_merge_begin(source, base_idx, aCount, bCount, aStart, bStart, diag, step_count):
+ first = ib.allocate("int64", (1,), name="first", scope="local")
+ mid = ib.allocate("int64", (1,), name="mid", scope="local")
+ last = ib.allocate("int64", (1,), name="last", scope="local")
+ first[0] = tvm.te.max(0, diag - bCount)
+ last[0] = tvm.te.min(diag, aCount)
+ with ib.while_loop(first[0] < last[0]):
+ mid[0] = (first[0] + last[0]) >> 1
+ a = source[base_idx + (aStart + mid[0])]
+ b = source[base_idx + (bStart + diag - 1 - mid[0])]
+ with ib.if_scope(compare(a, b)):
+ first[0] = mid[0] + 1
+ with ib.else_scope():
+ last[0] = mid[0]
+ return first, last
+
+ def serial_merge(
+ source,
+ dest,
+ source_idx,
+ dest_idx,
+ base_idx,
+ aCount,
+ bCount,
+ aStart,
+ bStart,
+ kStart,
+ diag,
+ step_count,
+ first,
+ last,
+ ):
+ i = ib.allocate("int64", (1,), name="i", scope="local")
+ j = ib.allocate("int64", (1,), name="j", scope="local")
+ i[0] = aStart + first[0]
+ j[0] = bStart + diag - last[0]
+ with ib.for_range(0, tvm.te.min(aCount + bCount - diag, step_count)) as count:
+ i_idx = base_idx + i[0]
+ j_idx = base_idx + j[0]
+ k_idx = base_idx + (kStart + diag + count)
+
+ def assign_i():
+ """assign i value to current output"""
+ dest[k_idx] = source[i_idx]
+ if values is not None:
+ dest_idx[k_idx] = source_idx[i_idx]
+ i[0] += 1
+
+ def assign_j():
+ """assign j value to current output"""
+ dest[k_idx] = source[j_idx]
+ if values is not None:
+ dest_idx[k_idx] = source_idx[j_idx]
+ j[0] += 1
+
+ ## if both of the iterators are in range
+ with ib.if_scope(tvm.tir.all(i[0] < aStart + aCount, j[0] < bStart + bCount)):
+ # compare them and insert whichever is next into the output
+ with ib.if_scope(compare(source[i_idx], source[j_idx])):
+ assign_i()
+ with ib.else_scope():
+ assign_j()
+ # otherwise, simply copy the remainder of the valid iterator to the output
+ with ib.else_scope():
+ with ib.if_scope(i[0] < aStart + aCount):
+ assign_i()
+ with ib.else_scope():
+ assign_j()
+
+ with ib.for_range(0, upper_lim - lower_lim, dtype="int64") as l2_width:
+ width = 2 << (l2_width + lower_lim)
# Define and launch the cuda kernel
with ib.new_scope():
- i = ib.allocate("int64", (1,), name="i", scope="local")
- j = ib.allocate("int64", (1,), name="j", scope="local")
- start = ib.allocate("int64", (1,), name="start", scope="local")
- middle = ib.allocate("int64", (1,), name="middle", scope="local")
- end = ib.allocate("int64", (1,), name="end", scope="local")
- tx = te.thread_axis("threadIdx.x")
- bx = te.thread_axis("blockIdx.x")
- ib.scope_attr(tx, "thread_extent", nthread_tx)
- # Reduce the number of blocks as the work per thread grows
- ib.scope_attr(
- bx,
- "thread_extent",
- tvm.tir.generic.cast(ceil_div(size, width * max_threads), "int32"),
- )
- tid = bx * nthread_tx + tx
Review comment:
Unfortunately, this slows me down by ~15 percent. :( Need to think about it more.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] merrymercy commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
merrymercy commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r590141913
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -96,6 +102,102 @@ def _sort_init(ib, shape, axis, keys_in, keys_out, values_out=None, value_init_f
return axis_mul_before, axis_mul_after
+## TODO(mbrookhart): These are effective optimziation hyperparametrs
+## Perhaps we can autotune?
+block_size = 128
+thread_work = 4
+
Review comment:
sure
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-793218495
@mbrookhart Can you comment on the difference in the way two-level merge path is done, between your implementation and `moderngpu` https://github.com/moderngpu/moderngpu/blob/master/src/moderngpu/kernel_mergesort.hxx#L97?
`moderngpu` launches a separate, device-wide merge path kernel, before invoking block level merge path kenrel. Your code seems to launch two merge path from one kernel, is that right?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r591921822
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -121,41 +221,85 @@ def _sort_common(
## On the final iteration of the algorithm, one thread will merge two sorted lists
## to sort the entire array
- max_threads = int(tvm.target.Target.current(allow_none=False).max_num_threads)
- nthread_tx = max_threads
- nthread_bx = ceil_div(size, max_threads)
- nthread_by = axis_mul_before
- nthread_bz = axis_mul_after
-
- def compare(a, b):
- """
- Compare a and b in proper ascending or descending order
- """
- if is_ascend:
- out = a <= b
- else:
- out = b <= a
- return out
-
- def bottom_up_merge(source, dest, source_idx, dest_idx, start, middle, end, even):
- """
- Merge the two sections of the array assigned to this thread
- """
- # pylint: disable=arguments-out-of-order
- # initialize iterators
- i = ib.allocate("int64", (1,), name="i", scope="local")
- j = ib.allocate("int64", (1,), name="j", scope="local")
- i[0] = start
- j[0] = middle
- # set up indexes
- base_idx = by * size * axis_mul_after + bz
- # iterate over the output loop
- with ib.for_range(0, end - start) as k:
- i_idx = base_idx + i[0] * axis_mul_after
- j_idx = base_idx + j[0] * axis_mul_after
- k_idx = base_idx + (k + start) * axis_mul_after
+ with ib.if_scope(size > 0):
Review comment:
maybe better to move this outside of `_sort_common`? To make indent not going too deep.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r589813103
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -136,93 +236,223 @@ def compare(a, b):
out = b <= a
return out
- def bottom_up_merge(source, dest, source_idx, dest_idx, start, middle, end, even):
- """
- Merge the two sections of the array assigned to this thread
- """
- # pylint: disable=arguments-out-of-order
- # initialize iterators
- i[0] = start
- j[0] = middle
- # set up indexes
- base_idx = by * size * axis_mul_after + bz
- # iterate over the output loop
- with ib.for_range(0, end - start) as k:
- i_idx = base_idx + i[0] * axis_mul_after
- j_idx = base_idx + j[0] * axis_mul_after
- k_idx = base_idx + (k + start) * axis_mul_after
-
- def swap_values(source, dest, source_idx, dest_idx):
- def assign_i():
- """assign i value to current output"""
- dest[k_idx] = source[i_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[i_idx]
- i[0] += 1
-
- def assign_j():
- """assign j value to current output"""
- dest[k_idx] = source[j_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[j_idx]
- j[0] += 1
-
- ## if both of the iterators are in range
- with ib.if_scope(tvm.tir.all(i[0] < middle, j[0] < end)):
- # compare them and insert whichever is next into the output
- with ib.if_scope(compare(source[i_idx], source[j_idx])):
- assign_i()
- with ib.else_scope():
- assign_j()
- # otherwise, simply copy the remainder of the valid iterator to the output
- with ib.else_scope():
- with ib.if_scope(i[0] < middle):
- assign_i()
- with ib.else_scope():
- assign_j()
+ # Sort the lower levels of the merge using odd-even sort, it's fast for small inputs
+ lower_lim = tvm.tir.generic.cast(
+ tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(block_size, "float64"))), "int64"
+ )
- # Switch which input is the source and which is the destination each iteration
- with ib.if_scope(even):
- swap_values(source, dest, source_idx, dest_idx)
- with ib.else_scope():
- swap_values(dest, source, dest_idx, source_idx)
-
- def mergesort(source, dest, source_idx, dest_idx, size, width, even):
- # calculate the start, mid, and end points of this section
- start[0] = width * tid
- with ib.if_scope(start[0] < size):
- middle[0] = tvm.te.min(start[0] + tvm.tir.indexdiv(width, 2), size)
- end[0] = tvm.te.min(start[0] + width, size)
- ## merge the start->middle and middle->end arrays
- bottom_up_merge(source, dest, source_idx, dest_idx, start[0], middle[0], end[0], even)
-
- lim = tvm.tir.generic.cast(
+ _odd_even_sort(
+ ib,
+ size,
+ axis_mul_before * axis_mul_after,
+ 1,
+ is_ascend,
+ keys,
+ keys_swap,
+ values,
+ values_swap,
+ )
+
+ upper_lim = tvm.tir.generic.cast(
tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(size, "float64"))), "int64"
)
- with ib.for_range(0, lim, dtype="int64") as l2_width:
- width = 2 << l2_width
+
+ def get_merge_begin(source, base_idx, aCount, bCount, aStart, bStart, diag, step_count):
+ first = ib.allocate("int64", (1,), name="first", scope="local")
+ mid = ib.allocate("int64", (1,), name="mid", scope="local")
+ last = ib.allocate("int64", (1,), name="last", scope="local")
+ first[0] = tvm.te.max(0, diag - bCount)
+ last[0] = tvm.te.min(diag, aCount)
+ with ib.while_loop(first[0] < last[0]):
+ mid[0] = (first[0] + last[0]) >> 1
+ a = source[base_idx + (aStart + mid[0])]
+ b = source[base_idx + (bStart + diag - 1 - mid[0])]
+ with ib.if_scope(compare(a, b)):
+ first[0] = mid[0] + 1
+ with ib.else_scope():
+ last[0] = mid[0]
+ return first, last
+
+ def serial_merge(
+ source,
+ dest,
+ source_idx,
+ dest_idx,
+ base_idx,
+ aCount,
+ bCount,
+ aStart,
+ bStart,
+ kStart,
+ diag,
+ step_count,
+ first,
+ last,
+ ):
+ i = ib.allocate("int64", (1,), name="i", scope="local")
+ j = ib.allocate("int64", (1,), name="j", scope="local")
+ i[0] = aStart + first[0]
+ j[0] = bStart + diag - last[0]
+ with ib.for_range(0, tvm.te.min(aCount + bCount - diag, step_count)) as count:
+ i_idx = base_idx + i[0]
+ j_idx = base_idx + j[0]
+ k_idx = base_idx + (kStart + diag + count)
+
+ def assign_i():
+ """assign i value to current output"""
+ dest[k_idx] = source[i_idx]
+ if values is not None:
+ dest_idx[k_idx] = source_idx[i_idx]
+ i[0] += 1
+
+ def assign_j():
+ """assign j value to current output"""
+ dest[k_idx] = source[j_idx]
+ if values is not None:
+ dest_idx[k_idx] = source_idx[j_idx]
+ j[0] += 1
+
+ ## if both of the iterators are in range
+ with ib.if_scope(tvm.tir.all(i[0] < aStart + aCount, j[0] < bStart + bCount)):
+ # compare them and insert whichever is next into the output
+ with ib.if_scope(compare(source[i_idx], source[j_idx])):
+ assign_i()
+ with ib.else_scope():
+ assign_j()
+ # otherwise, simply copy the remainder of the valid iterator to the output
+ with ib.else_scope():
+ with ib.if_scope(i[0] < aStart + aCount):
+ assign_i()
+ with ib.else_scope():
+ assign_j()
+
+ with ib.for_range(0, upper_lim - lower_lim, dtype="int64") as l2_width:
+ width = 2 << (l2_width + lower_lim)
# Define and launch the cuda kernel
with ib.new_scope():
- i = ib.allocate("int64", (1,), name="i", scope="local")
- j = ib.allocate("int64", (1,), name="j", scope="local")
- start = ib.allocate("int64", (1,), name="start", scope="local")
- middle = ib.allocate("int64", (1,), name="middle", scope="local")
- end = ib.allocate("int64", (1,), name="end", scope="local")
- tx = te.thread_axis("threadIdx.x")
- bx = te.thread_axis("blockIdx.x")
- ib.scope_attr(tx, "thread_extent", nthread_tx)
- # Reduce the number of blocks as the work per thread grows
- ib.scope_attr(
- bx,
- "thread_extent",
- tvm.tir.generic.cast(ceil_div(size, width * max_threads), "int32"),
- )
- tid = bx * nthread_tx + tx
Review comment:
Yes I think that would work. I can help you testing on amd
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-796063236
Interesting! I did a binary search for which size fails on VK, size below 2048 works fine, but starting from 2049 it fails.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] mbrookhart commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r589805036
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -136,93 +236,223 @@ def compare(a, b):
out = b <= a
return out
- def bottom_up_merge(source, dest, source_idx, dest_idx, start, middle, end, even):
- """
- Merge the two sections of the array assigned to this thread
- """
- # pylint: disable=arguments-out-of-order
- # initialize iterators
- i[0] = start
- j[0] = middle
- # set up indexes
- base_idx = by * size * axis_mul_after + bz
- # iterate over the output loop
- with ib.for_range(0, end - start) as k:
- i_idx = base_idx + i[0] * axis_mul_after
- j_idx = base_idx + j[0] * axis_mul_after
- k_idx = base_idx + (k + start) * axis_mul_after
-
- def swap_values(source, dest, source_idx, dest_idx):
- def assign_i():
- """assign i value to current output"""
- dest[k_idx] = source[i_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[i_idx]
- i[0] += 1
-
- def assign_j():
- """assign j value to current output"""
- dest[k_idx] = source[j_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[j_idx]
- j[0] += 1
-
- ## if both of the iterators are in range
- with ib.if_scope(tvm.tir.all(i[0] < middle, j[0] < end)):
- # compare them and insert whichever is next into the output
- with ib.if_scope(compare(source[i_idx], source[j_idx])):
- assign_i()
- with ib.else_scope():
- assign_j()
- # otherwise, simply copy the remainder of the valid iterator to the output
- with ib.else_scope():
- with ib.if_scope(i[0] < middle):
- assign_i()
- with ib.else_scope():
- assign_j()
+ # Sort the lower levels of the merge using odd-even sort, it's fast for small inputs
+ lower_lim = tvm.tir.generic.cast(
+ tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(block_size, "float64"))), "int64"
+ )
- # Switch which input is the source and which is the destination each iteration
- with ib.if_scope(even):
- swap_values(source, dest, source_idx, dest_idx)
- with ib.else_scope():
- swap_values(dest, source, dest_idx, source_idx)
-
- def mergesort(source, dest, source_idx, dest_idx, size, width, even):
- # calculate the start, mid, and end points of this section
- start[0] = width * tid
- with ib.if_scope(start[0] < size):
- middle[0] = tvm.te.min(start[0] + tvm.tir.indexdiv(width, 2), size)
- end[0] = tvm.te.min(start[0] + width, size)
- ## merge the start->middle and middle->end arrays
- bottom_up_merge(source, dest, source_idx, dest_idx, start[0], middle[0], end[0], even)
-
- lim = tvm.tir.generic.cast(
+ _odd_even_sort(
+ ib,
+ size,
+ axis_mul_before * axis_mul_after,
+ 1,
+ is_ascend,
+ keys,
+ keys_swap,
+ values,
+ values_swap,
+ )
+
+ upper_lim = tvm.tir.generic.cast(
tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(size, "float64"))), "int64"
)
- with ib.for_range(0, lim, dtype="int64") as l2_width:
- width = 2 << l2_width
+
+ def get_merge_begin(source, base_idx, aCount, bCount, aStart, bStart, diag, step_count):
+ first = ib.allocate("int64", (1,), name="first", scope="local")
+ mid = ib.allocate("int64", (1,), name="mid", scope="local")
+ last = ib.allocate("int64", (1,), name="last", scope="local")
+ first[0] = tvm.te.max(0, diag - bCount)
+ last[0] = tvm.te.min(diag, aCount)
+ with ib.while_loop(first[0] < last[0]):
+ mid[0] = (first[0] + last[0]) >> 1
+ a = source[base_idx + (aStart + mid[0])]
+ b = source[base_idx + (bStart + diag - 1 - mid[0])]
+ with ib.if_scope(compare(a, b)):
+ first[0] = mid[0] + 1
+ with ib.else_scope():
+ last[0] = mid[0]
+ return first, last
+
+ def serial_merge(
+ source,
+ dest,
+ source_idx,
+ dest_idx,
+ base_idx,
+ aCount,
+ bCount,
+ aStart,
+ bStart,
+ kStart,
+ diag,
+ step_count,
+ first,
+ last,
+ ):
+ i = ib.allocate("int64", (1,), name="i", scope="local")
+ j = ib.allocate("int64", (1,), name="j", scope="local")
+ i[0] = aStart + first[0]
+ j[0] = bStart + diag - last[0]
+ with ib.for_range(0, tvm.te.min(aCount + bCount - diag, step_count)) as count:
+ i_idx = base_idx + i[0]
+ j_idx = base_idx + j[0]
+ k_idx = base_idx + (kStart + diag + count)
+
+ def assign_i():
+ """assign i value to current output"""
+ dest[k_idx] = source[i_idx]
+ if values is not None:
+ dest_idx[k_idx] = source_idx[i_idx]
+ i[0] += 1
+
+ def assign_j():
+ """assign j value to current output"""
+ dest[k_idx] = source[j_idx]
+ if values is not None:
+ dest_idx[k_idx] = source_idx[j_idx]
+ j[0] += 1
+
+ ## if both of the iterators are in range
+ with ib.if_scope(tvm.tir.all(i[0] < aStart + aCount, j[0] < bStart + bCount)):
+ # compare them and insert whichever is next into the output
+ with ib.if_scope(compare(source[i_idx], source[j_idx])):
+ assign_i()
+ with ib.else_scope():
+ assign_j()
+ # otherwise, simply copy the remainder of the valid iterator to the output
+ with ib.else_scope():
+ with ib.if_scope(i[0] < aStart + aCount):
+ assign_i()
+ with ib.else_scope():
+ assign_j()
+
+ with ib.for_range(0, upper_lim - lower_lim, dtype="int64") as l2_width:
+ width = 2 << (l2_width + lower_lim)
# Define and launch the cuda kernel
with ib.new_scope():
- i = ib.allocate("int64", (1,), name="i", scope="local")
- j = ib.allocate("int64", (1,), name="j", scope="local")
- start = ib.allocate("int64", (1,), name="start", scope="local")
- middle = ib.allocate("int64", (1,), name="middle", scope="local")
- end = ib.allocate("int64", (1,), name="end", scope="local")
- tx = te.thread_axis("threadIdx.x")
- bx = te.thread_axis("blockIdx.x")
- ib.scope_attr(tx, "thread_extent", nthread_tx)
- # Reduce the number of blocks as the work per thread grows
- ib.scope_attr(
- bx,
- "thread_extent",
- tvm.tir.generic.cast(ceil_div(size, width * max_threads), "int32"),
- )
- tid = bx * nthread_tx + tx
Review comment:
@masahi I'm kind of surprised spirv doesn't like it, I was doing a thread calculation before. Perhaps it's just complaining about nthread_tx? Perhaps we just make `ntx = tvm.tir.generic.cast(tvm.te.min(max_threads, width), "int32")` into `ntx = max_threads`. That should only hit performance slightly and on a couple of levels (there are already if `tid < x` statements everywhere), would that work for spirv?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] mbrookhart commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r590536212
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -96,6 +102,102 @@ def _sort_init(ib, shape, axis, keys_in, keys_out, values_out=None, value_init_f
return axis_mul_before, axis_mul_after
+## TODO(mbrookhart): These are effective optimziation hyperparametrs
+## Perhaps we can autotune?
+block_size = 128
+thread_work = 4
+
Review comment:
Any tips on how to do 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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] mbrookhart commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r590506226
##########
File path: src/tir/transforms/storage_access.cc
##########
@@ -132,6 +132,7 @@ void StorageAccessVisitor::VisitStmt_(const AttrStmtNode* op) {
StmtExprVisitor::VisitStmt_(op);
}
env_threads_.pop_back();
+ } else if (op->attr_key == attr::hand_threaded) {
} else {
Review comment:
I'll add a comment, but the point of this is to skip thread planning for the enclosed block, i.e., not to visit the statements in this block. I do that because 1) it's already threaded and 2) the thread planning pass can't handle the per-thread read/write/sync I'm doing.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r589766328
##########
File path: tests/python/relay/test_op_level6.py
##########
@@ -88,7 +88,8 @@ def verify_argsort(shape, axis, is_ascend, dtype, is_dyn=False):
verify_argsort((2, 3, 4), axis=0, is_ascend=False, dtype=dtype, is_dyn=is_dyn)
verify_argsort((1, 4, 6), axis=1, is_ascend=True, dtype=dtype, is_dyn=is_dyn)
verify_argsort((3, 5, 6), axis=-1, is_ascend=False, dtype=dtype, is_dyn=is_dyn)
- verify_argsort((3, 2000, 6), axis=1, is_ascend=False, dtype=dtype, is_dyn=is_dyn)
+ verify_argsort((3, 6000, 6), axis=1, is_ascend=False, dtype=dtype, is_dyn=is_dyn)
+ verify_argsort((12000, 8, 8), axis=0, is_ascend=True, dtype=dtype, is_dyn=is_dyn)
verify_argsort((1, 122640), axis=1, is_ascend=False, dtype=dtype, is_dyn=is_dyn)
Review comment:
Can we add a workload from your benchmark script, that was causing an issue before?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] mbrookhart commented on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-794195458
@masahi I was having a hard time getting partition information back out when I used a seperate kernel, mostly due to handling memory in IR builder. Running the binary search in the kernel that will eventual merge the partitions was easier to handle from a memory-passing sense. I could go back and try temporary allocations to see if the added complexity gets us any performance improvement. Still thinking about the threading for vulkan.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r589855001
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -136,93 +236,223 @@ def compare(a, b):
out = b <= a
return out
- def bottom_up_merge(source, dest, source_idx, dest_idx, start, middle, end, even):
- """
- Merge the two sections of the array assigned to this thread
- """
- # pylint: disable=arguments-out-of-order
- # initialize iterators
- i[0] = start
- j[0] = middle
- # set up indexes
- base_idx = by * size * axis_mul_after + bz
- # iterate over the output loop
- with ib.for_range(0, end - start) as k:
- i_idx = base_idx + i[0] * axis_mul_after
- j_idx = base_idx + j[0] * axis_mul_after
- k_idx = base_idx + (k + start) * axis_mul_after
-
- def swap_values(source, dest, source_idx, dest_idx):
- def assign_i():
- """assign i value to current output"""
- dest[k_idx] = source[i_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[i_idx]
- i[0] += 1
-
- def assign_j():
- """assign j value to current output"""
- dest[k_idx] = source[j_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[j_idx]
- j[0] += 1
-
- ## if both of the iterators are in range
- with ib.if_scope(tvm.tir.all(i[0] < middle, j[0] < end)):
- # compare them and insert whichever is next into the output
- with ib.if_scope(compare(source[i_idx], source[j_idx])):
- assign_i()
- with ib.else_scope():
- assign_j()
- # otherwise, simply copy the remainder of the valid iterator to the output
- with ib.else_scope():
- with ib.if_scope(i[0] < middle):
- assign_i()
- with ib.else_scope():
- assign_j()
+ # Sort the lower levels of the merge using odd-even sort, it's fast for small inputs
+ lower_lim = tvm.tir.generic.cast(
+ tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(block_size, "float64"))), "int64"
+ )
- # Switch which input is the source and which is the destination each iteration
- with ib.if_scope(even):
- swap_values(source, dest, source_idx, dest_idx)
- with ib.else_scope():
- swap_values(dest, source, dest_idx, source_idx)
-
- def mergesort(source, dest, source_idx, dest_idx, size, width, even):
- # calculate the start, mid, and end points of this section
- start[0] = width * tid
- with ib.if_scope(start[0] < size):
- middle[0] = tvm.te.min(start[0] + tvm.tir.indexdiv(width, 2), size)
- end[0] = tvm.te.min(start[0] + width, size)
- ## merge the start->middle and middle->end arrays
- bottom_up_merge(source, dest, source_idx, dest_idx, start[0], middle[0], end[0], even)
-
- lim = tvm.tir.generic.cast(
+ _odd_even_sort(
+ ib,
+ size,
+ axis_mul_before * axis_mul_after,
+ 1,
+ is_ascend,
+ keys,
+ keys_swap,
+ values,
+ values_swap,
+ )
+
+ upper_lim = tvm.tir.generic.cast(
tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(size, "float64"))), "int64"
)
- with ib.for_range(0, lim, dtype="int64") as l2_width:
- width = 2 << l2_width
+
+ def get_merge_begin(source, base_idx, aCount, bCount, aStart, bStart, diag, step_count):
+ first = ib.allocate("int64", (1,), name="first", scope="local")
+ mid = ib.allocate("int64", (1,), name="mid", scope="local")
+ last = ib.allocate("int64", (1,), name="last", scope="local")
+ first[0] = tvm.te.max(0, diag - bCount)
+ last[0] = tvm.te.min(diag, aCount)
+ with ib.while_loop(first[0] < last[0]):
+ mid[0] = (first[0] + last[0]) >> 1
Review comment:
Do we need to allocate `mid`? Since you don't need to update this variable.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] mbrookhart commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r590686714
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -136,93 +236,223 @@ def compare(a, b):
out = b <= a
return out
- def bottom_up_merge(source, dest, source_idx, dest_idx, start, middle, end, even):
- """
- Merge the two sections of the array assigned to this thread
- """
- # pylint: disable=arguments-out-of-order
- # initialize iterators
- i[0] = start
- j[0] = middle
- # set up indexes
- base_idx = by * size * axis_mul_after + bz
- # iterate over the output loop
- with ib.for_range(0, end - start) as k:
- i_idx = base_idx + i[0] * axis_mul_after
- j_idx = base_idx + j[0] * axis_mul_after
- k_idx = base_idx + (k + start) * axis_mul_after
-
- def swap_values(source, dest, source_idx, dest_idx):
- def assign_i():
- """assign i value to current output"""
- dest[k_idx] = source[i_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[i_idx]
- i[0] += 1
-
- def assign_j():
- """assign j value to current output"""
- dest[k_idx] = source[j_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[j_idx]
- j[0] += 1
-
- ## if both of the iterators are in range
- with ib.if_scope(tvm.tir.all(i[0] < middle, j[0] < end)):
- # compare them and insert whichever is next into the output
- with ib.if_scope(compare(source[i_idx], source[j_idx])):
- assign_i()
- with ib.else_scope():
- assign_j()
- # otherwise, simply copy the remainder of the valid iterator to the output
- with ib.else_scope():
- with ib.if_scope(i[0] < middle):
- assign_i()
- with ib.else_scope():
- assign_j()
+ # Sort the lower levels of the merge using odd-even sort, it's fast for small inputs
+ lower_lim = tvm.tir.generic.cast(
+ tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(block_size, "float64"))), "int64"
+ )
- # Switch which input is the source and which is the destination each iteration
- with ib.if_scope(even):
- swap_values(source, dest, source_idx, dest_idx)
- with ib.else_scope():
- swap_values(dest, source, dest_idx, source_idx)
-
- def mergesort(source, dest, source_idx, dest_idx, size, width, even):
- # calculate the start, mid, and end points of this section
- start[0] = width * tid
- with ib.if_scope(start[0] < size):
- middle[0] = tvm.te.min(start[0] + tvm.tir.indexdiv(width, 2), size)
- end[0] = tvm.te.min(start[0] + width, size)
- ## merge the start->middle and middle->end arrays
- bottom_up_merge(source, dest, source_idx, dest_idx, start[0], middle[0], end[0], even)
-
- lim = tvm.tir.generic.cast(
+ _odd_even_sort(
+ ib,
+ size,
+ axis_mul_before * axis_mul_after,
+ 1,
+ is_ascend,
+ keys,
+ keys_swap,
+ values,
+ values_swap,
+ )
+
+ upper_lim = tvm.tir.generic.cast(
tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(size, "float64"))), "int64"
)
- with ib.for_range(0, lim, dtype="int64") as l2_width:
- width = 2 << l2_width
+
+ def get_merge_begin(source, base_idx, aCount, bCount, aStart, bStart, diag, step_count):
+ first = ib.allocate("int64", (1,), name="first", scope="local")
+ mid = ib.allocate("int64", (1,), name="mid", scope="local")
+ last = ib.allocate("int64", (1,), name="last", scope="local")
+ first[0] = tvm.te.max(0, diag - bCount)
+ last[0] = tvm.te.min(diag, aCount)
+ with ib.while_loop(first[0] < last[0]):
+ mid[0] = (first[0] + last[0]) >> 1
+ a = source[base_idx + (aStart + mid[0])]
+ b = source[base_idx + (bStart + diag - 1 - mid[0])]
+ with ib.if_scope(compare(a, b)):
+ first[0] = mid[0] + 1
+ with ib.else_scope():
+ last[0] = mid[0]
+ return first, last
+
+ def serial_merge(
+ source,
+ dest,
+ source_idx,
+ dest_idx,
+ base_idx,
+ aCount,
+ bCount,
+ aStart,
+ bStart,
+ kStart,
+ diag,
+ step_count,
+ first,
+ last,
+ ):
+ i = ib.allocate("int64", (1,), name="i", scope="local")
+ j = ib.allocate("int64", (1,), name="j", scope="local")
+ i[0] = aStart + first[0]
+ j[0] = bStart + diag - last[0]
+ with ib.for_range(0, tvm.te.min(aCount + bCount - diag, step_count)) as count:
+ i_idx = base_idx + i[0]
+ j_idx = base_idx + j[0]
+ k_idx = base_idx + (kStart + diag + count)
+
+ def assign_i():
+ """assign i value to current output"""
+ dest[k_idx] = source[i_idx]
+ if values is not None:
+ dest_idx[k_idx] = source_idx[i_idx]
+ i[0] += 1
+
+ def assign_j():
+ """assign j value to current output"""
+ dest[k_idx] = source[j_idx]
+ if values is not None:
+ dest_idx[k_idx] = source_idx[j_idx]
+ j[0] += 1
+
+ ## if both of the iterators are in range
+ with ib.if_scope(tvm.tir.all(i[0] < aStart + aCount, j[0] < bStart + bCount)):
+ # compare them and insert whichever is next into the output
+ with ib.if_scope(compare(source[i_idx], source[j_idx])):
+ assign_i()
+ with ib.else_scope():
+ assign_j()
+ # otherwise, simply copy the remainder of the valid iterator to the output
+ with ib.else_scope():
+ with ib.if_scope(i[0] < aStart + aCount):
+ assign_i()
+ with ib.else_scope():
+ assign_j()
+
+ with ib.for_range(0, upper_lim - lower_lim, dtype="int64") as l2_width:
+ width = 2 << (l2_width + lower_lim)
# Define and launch the cuda kernel
with ib.new_scope():
- i = ib.allocate("int64", (1,), name="i", scope="local")
- j = ib.allocate("int64", (1,), name="j", scope="local")
- start = ib.allocate("int64", (1,), name="start", scope="local")
- middle = ib.allocate("int64", (1,), name="middle", scope="local")
- end = ib.allocate("int64", (1,), name="end", scope="local")
- tx = te.thread_axis("threadIdx.x")
- bx = te.thread_axis("blockIdx.x")
- ib.scope_attr(tx, "thread_extent", nthread_tx)
- # Reduce the number of blocks as the work per thread grows
- ib.scope_attr(
- bx,
- "thread_extent",
- tvm.tir.generic.cast(ceil_div(size, width * max_threads), "int32"),
- )
- tid = bx * nthread_tx + tx
Review comment:
Yeah, slowdown on cuda from making this change. Still correct, just slower. I'm not sure how to pass target information into the irbuilder to do different thinks for different backends.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] mbrookhart commented on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-796032638
I validated that that threading change passed tests when I applied it to cuda, it just hurt performance :( I wonder if I'm hitting another spirv codegen issue?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r589764976
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -96,6 +102,102 @@ def _sort_init(ib, shape, axis, keys_in, keys_out, values_out=None, value_init_f
return axis_mul_before, axis_mul_after
+## TODO(mbrookhart): These are effective optimziation hyperparametrs
+## Perhaps we can autotune?
+block_size = 128
+thread_work = 4
+
Review comment:
cc @merrymercy can we use autotvm for this? i.e. a traditional auto tuner way?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] mbrookhart commented on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-796183135
@masahi that might be the point at which I switch from single level mergepath to dual level mergepath. I wonder if vulkan is having trouble with the extra control flow there.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] mbrookhart commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r590506792
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -136,93 +236,223 @@ def compare(a, b):
out = b <= a
return out
- def bottom_up_merge(source, dest, source_idx, dest_idx, start, middle, end, even):
- """
- Merge the two sections of the array assigned to this thread
- """
- # pylint: disable=arguments-out-of-order
- # initialize iterators
- i[0] = start
- j[0] = middle
- # set up indexes
- base_idx = by * size * axis_mul_after + bz
- # iterate over the output loop
- with ib.for_range(0, end - start) as k:
- i_idx = base_idx + i[0] * axis_mul_after
- j_idx = base_idx + j[0] * axis_mul_after
- k_idx = base_idx + (k + start) * axis_mul_after
-
- def swap_values(source, dest, source_idx, dest_idx):
- def assign_i():
- """assign i value to current output"""
- dest[k_idx] = source[i_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[i_idx]
- i[0] += 1
-
- def assign_j():
- """assign j value to current output"""
- dest[k_idx] = source[j_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[j_idx]
- j[0] += 1
-
- ## if both of the iterators are in range
- with ib.if_scope(tvm.tir.all(i[0] < middle, j[0] < end)):
- # compare them and insert whichever is next into the output
- with ib.if_scope(compare(source[i_idx], source[j_idx])):
- assign_i()
- with ib.else_scope():
- assign_j()
- # otherwise, simply copy the remainder of the valid iterator to the output
- with ib.else_scope():
- with ib.if_scope(i[0] < middle):
- assign_i()
- with ib.else_scope():
- assign_j()
+ # Sort the lower levels of the merge using odd-even sort, it's fast for small inputs
+ lower_lim = tvm.tir.generic.cast(
+ tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(block_size, "float64"))), "int64"
+ )
- # Switch which input is the source and which is the destination each iteration
- with ib.if_scope(even):
- swap_values(source, dest, source_idx, dest_idx)
- with ib.else_scope():
- swap_values(dest, source, dest_idx, source_idx)
-
- def mergesort(source, dest, source_idx, dest_idx, size, width, even):
- # calculate the start, mid, and end points of this section
- start[0] = width * tid
- with ib.if_scope(start[0] < size):
- middle[0] = tvm.te.min(start[0] + tvm.tir.indexdiv(width, 2), size)
- end[0] = tvm.te.min(start[0] + width, size)
- ## merge the start->middle and middle->end arrays
- bottom_up_merge(source, dest, source_idx, dest_idx, start[0], middle[0], end[0], even)
-
- lim = tvm.tir.generic.cast(
+ _odd_even_sort(
+ ib,
+ size,
+ axis_mul_before * axis_mul_after,
+ 1,
+ is_ascend,
+ keys,
+ keys_swap,
+ values,
+ values_swap,
+ )
+
+ upper_lim = tvm.tir.generic.cast(
tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(size, "float64"))), "int64"
)
- with ib.for_range(0, lim, dtype="int64") as l2_width:
- width = 2 << l2_width
+
+ def get_merge_begin(source, base_idx, aCount, bCount, aStart, bStart, diag, step_count):
+ first = ib.allocate("int64", (1,), name="first", scope="local")
+ mid = ib.allocate("int64", (1,), name="mid", scope="local")
+ last = ib.allocate("int64", (1,), name="last", scope="local")
+ first[0] = tvm.te.max(0, diag - bCount)
+ last[0] = tvm.te.min(diag, aCount)
+ with ib.while_loop(first[0] < last[0]):
+ mid[0] = (first[0] + last[0]) >> 1
+ a = source[base_idx + (aStart + mid[0])]
+ b = source[base_idx + (bStart + diag - 1 - mid[0])]
+ with ib.if_scope(compare(a, b)):
+ first[0] = mid[0] + 1
+ with ib.else_scope():
+ last[0] = mid[0]
+ return first, last
Review comment:
I'm not sure it matters because this is IR builder, everything in this function ends up getting inlined, but I'll take another look at the style. :+1:
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-794489881
There is an issue with dynamic topk? I saw the same failure yesterday
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-796070741
@mbrookhart Also on the failure from `test_any.py`: The last test case tests on 0 size input, and it fails both rocm and vulkan. On VK I get segfault, on rocm I get `Floating point exception (core dumped)`. Maybe you want to check if your code is safe for empty input?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-799954739
Thanks @mbrookhart
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r589857594
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -136,93 +236,223 @@ def compare(a, b):
out = b <= a
return out
- def bottom_up_merge(source, dest, source_idx, dest_idx, start, middle, end, even):
- """
- Merge the two sections of the array assigned to this thread
- """
- # pylint: disable=arguments-out-of-order
- # initialize iterators
- i[0] = start
- j[0] = middle
- # set up indexes
- base_idx = by * size * axis_mul_after + bz
- # iterate over the output loop
- with ib.for_range(0, end - start) as k:
- i_idx = base_idx + i[0] * axis_mul_after
- j_idx = base_idx + j[0] * axis_mul_after
- k_idx = base_idx + (k + start) * axis_mul_after
-
- def swap_values(source, dest, source_idx, dest_idx):
- def assign_i():
- """assign i value to current output"""
- dest[k_idx] = source[i_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[i_idx]
- i[0] += 1
-
- def assign_j():
- """assign j value to current output"""
- dest[k_idx] = source[j_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[j_idx]
- j[0] += 1
-
- ## if both of the iterators are in range
- with ib.if_scope(tvm.tir.all(i[0] < middle, j[0] < end)):
- # compare them and insert whichever is next into the output
- with ib.if_scope(compare(source[i_idx], source[j_idx])):
- assign_i()
- with ib.else_scope():
- assign_j()
- # otherwise, simply copy the remainder of the valid iterator to the output
- with ib.else_scope():
- with ib.if_scope(i[0] < middle):
- assign_i()
- with ib.else_scope():
- assign_j()
+ # Sort the lower levels of the merge using odd-even sort, it's fast for small inputs
+ lower_lim = tvm.tir.generic.cast(
+ tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(block_size, "float64"))), "int64"
+ )
- # Switch which input is the source and which is the destination each iteration
- with ib.if_scope(even):
- swap_values(source, dest, source_idx, dest_idx)
- with ib.else_scope():
- swap_values(dest, source, dest_idx, source_idx)
-
- def mergesort(source, dest, source_idx, dest_idx, size, width, even):
- # calculate the start, mid, and end points of this section
- start[0] = width * tid
- with ib.if_scope(start[0] < size):
- middle[0] = tvm.te.min(start[0] + tvm.tir.indexdiv(width, 2), size)
- end[0] = tvm.te.min(start[0] + width, size)
- ## merge the start->middle and middle->end arrays
- bottom_up_merge(source, dest, source_idx, dest_idx, start[0], middle[0], end[0], even)
-
- lim = tvm.tir.generic.cast(
+ _odd_even_sort(
+ ib,
+ size,
+ axis_mul_before * axis_mul_after,
+ 1,
+ is_ascend,
+ keys,
+ keys_swap,
+ values,
+ values_swap,
+ )
+
+ upper_lim = tvm.tir.generic.cast(
tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(size, "float64"))), "int64"
)
- with ib.for_range(0, lim, dtype="int64") as l2_width:
- width = 2 << l2_width
+
+ def get_merge_begin(source, base_idx, aCount, bCount, aStart, bStart, diag, step_count):
+ first = ib.allocate("int64", (1,), name="first", scope="local")
+ mid = ib.allocate("int64", (1,), name="mid", scope="local")
+ last = ib.allocate("int64", (1,), name="last", scope="local")
+ first[0] = tvm.te.max(0, diag - bCount)
+ last[0] = tvm.te.min(diag, aCount)
+ with ib.while_loop(first[0] < last[0]):
+ mid[0] = (first[0] + last[0]) >> 1
+ a = source[base_idx + (aStart + mid[0])]
+ b = source[base_idx + (bStart + diag - 1 - mid[0])]
+ with ib.if_scope(compare(a, b)):
+ first[0] = mid[0] + 1
+ with ib.else_scope():
+ last[0] = mid[0]
+ return first, last
Review comment:
How about returning `first[0]` and `last[0]`? I'm not comfortable with returning local allocations.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tkonolige commented on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
tkonolige commented on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-793093494
Can we just run the full thread block size and nop on the extra threads?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r589858178
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -136,93 +236,223 @@ def compare(a, b):
out = b <= a
return out
- def bottom_up_merge(source, dest, source_idx, dest_idx, start, middle, end, even):
- """
- Merge the two sections of the array assigned to this thread
- """
- # pylint: disable=arguments-out-of-order
- # initialize iterators
- i[0] = start
- j[0] = middle
- # set up indexes
- base_idx = by * size * axis_mul_after + bz
- # iterate over the output loop
- with ib.for_range(0, end - start) as k:
- i_idx = base_idx + i[0] * axis_mul_after
- j_idx = base_idx + j[0] * axis_mul_after
- k_idx = base_idx + (k + start) * axis_mul_after
-
- def swap_values(source, dest, source_idx, dest_idx):
- def assign_i():
- """assign i value to current output"""
- dest[k_idx] = source[i_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[i_idx]
- i[0] += 1
-
- def assign_j():
- """assign j value to current output"""
- dest[k_idx] = source[j_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[j_idx]
- j[0] += 1
-
- ## if both of the iterators are in range
- with ib.if_scope(tvm.tir.all(i[0] < middle, j[0] < end)):
- # compare them and insert whichever is next into the output
- with ib.if_scope(compare(source[i_idx], source[j_idx])):
- assign_i()
- with ib.else_scope():
- assign_j()
- # otherwise, simply copy the remainder of the valid iterator to the output
- with ib.else_scope():
- with ib.if_scope(i[0] < middle):
- assign_i()
- with ib.else_scope():
- assign_j()
+ # Sort the lower levels of the merge using odd-even sort, it's fast for small inputs
+ lower_lim = tvm.tir.generic.cast(
+ tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(block_size, "float64"))), "int64"
+ )
- # Switch which input is the source and which is the destination each iteration
- with ib.if_scope(even):
- swap_values(source, dest, source_idx, dest_idx)
- with ib.else_scope():
- swap_values(dest, source, dest_idx, source_idx)
-
- def mergesort(source, dest, source_idx, dest_idx, size, width, even):
- # calculate the start, mid, and end points of this section
- start[0] = width * tid
- with ib.if_scope(start[0] < size):
- middle[0] = tvm.te.min(start[0] + tvm.tir.indexdiv(width, 2), size)
- end[0] = tvm.te.min(start[0] + width, size)
- ## merge the start->middle and middle->end arrays
- bottom_up_merge(source, dest, source_idx, dest_idx, start[0], middle[0], end[0], even)
-
- lim = tvm.tir.generic.cast(
+ _odd_even_sort(
+ ib,
+ size,
+ axis_mul_before * axis_mul_after,
+ 1,
+ is_ascend,
+ keys,
+ keys_swap,
+ values,
+ values_swap,
+ )
+
+ upper_lim = tvm.tir.generic.cast(
tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(size, "float64"))), "int64"
)
- with ib.for_range(0, lim, dtype="int64") as l2_width:
- width = 2 << l2_width
+
+ def get_merge_begin(source, base_idx, aCount, bCount, aStart, bStart, diag, step_count):
+ first = ib.allocate("int64", (1,), name="first", scope="local")
+ mid = ib.allocate("int64", (1,), name="mid", scope="local")
+ last = ib.allocate("int64", (1,), name="last", scope="local")
+ first[0] = tvm.te.max(0, diag - bCount)
+ last[0] = tvm.te.min(diag, aCount)
+ with ib.while_loop(first[0] < last[0]):
+ mid[0] = (first[0] + last[0]) >> 1
+ a = source[base_idx + (aStart + mid[0])]
+ b = source[base_idx + (bStart + diag - 1 - mid[0])]
+ with ib.if_scope(compare(a, b)):
+ first[0] = mid[0] + 1
+ with ib.else_scope():
+ last[0] = mid[0]
+ return first, last
+
+ def serial_merge(
+ source,
+ dest,
+ source_idx,
+ dest_idx,
+ base_idx,
+ aCount,
+ bCount,
+ aStart,
+ bStart,
+ kStart,
+ diag,
+ step_count,
+ first,
+ last,
+ ):
+ i = ib.allocate("int64", (1,), name="i", scope="local")
+ j = ib.allocate("int64", (1,), name="j", scope="local")
+ i[0] = aStart + first[0]
+ j[0] = bStart + diag - last[0]
+ with ib.for_range(0, tvm.te.min(aCount + bCount - diag, step_count)) as count:
+ i_idx = base_idx + i[0]
+ j_idx = base_idx + j[0]
+ k_idx = base_idx + (kStart + diag + count)
+
+ def assign_i():
+ """assign i value to current output"""
+ dest[k_idx] = source[i_idx]
+ if values is not None:
+ dest_idx[k_idx] = source_idx[i_idx]
+ i[0] += 1
+
+ def assign_j():
+ """assign j value to current output"""
+ dest[k_idx] = source[j_idx]
+ if values is not None:
+ dest_idx[k_idx] = source_idx[j_idx]
+ j[0] += 1
+
+ ## if both of the iterators are in range
+ with ib.if_scope(tvm.tir.all(i[0] < aStart + aCount, j[0] < bStart + bCount)):
+ # compare them and insert whichever is next into the output
+ with ib.if_scope(compare(source[i_idx], source[j_idx])):
+ assign_i()
+ with ib.else_scope():
+ assign_j()
+ # otherwise, simply copy the remainder of the valid iterator to the output
+ with ib.else_scope():
+ with ib.if_scope(i[0] < aStart + aCount):
+ assign_i()
+ with ib.else_scope():
+ assign_j()
+
+ with ib.for_range(0, upper_lim - lower_lim, dtype="int64") as l2_width:
+ width = 2 << (l2_width + lower_lim)
# Define and launch the cuda kernel
with ib.new_scope():
- i = ib.allocate("int64", (1,), name="i", scope="local")
- j = ib.allocate("int64", (1,), name="j", scope="local")
- start = ib.allocate("int64", (1,), name="start", scope="local")
- middle = ib.allocate("int64", (1,), name="middle", scope="local")
- end = ib.allocate("int64", (1,), name="end", scope="local")
- tx = te.thread_axis("threadIdx.x")
- bx = te.thread_axis("blockIdx.x")
- ib.scope_attr(tx, "thread_extent", nthread_tx)
- # Reduce the number of blocks as the work per thread grows
- ib.scope_attr(
- bx,
- "thread_extent",
- tvm.tir.generic.cast(ceil_div(size, width * max_threads), "int32"),
- )
- tid = bx * nthread_tx + tx
-
- by = te.thread_axis("blockIdx.y")
- bz = te.thread_axis("blockIdx.z")
- ib.scope_attr(by, "thread_extent", nthread_by)
- ib.scope_attr(bz, "thread_extent", nthread_bz)
+ ntx = tvm.tir.generic.cast(tvm.te.min(max_threads, width), "int32")
+ nbx = tvm.tir.generic.cast(ceil_div(width, max_threads * thread_work), "int32")
+ nbz = tvm.tir.generic.cast(ceil_div(size, width), "int32")
+ tx, bx, by, bz = _get_threads(ib, ntx, nbx, nthread_by, nbz)
+
+ def mergepath(
+ source,
+ dest,
+ source_idx,
+ dest_idx,
+ aCount,
+ bCount,
+ aStart,
+ bStart,
+ kStart,
+ step_count,
+ even,
+ ):
+ # pylint: disable=arguments-out-of-order
+ def merge(source, dest, source_idx, dest_idx):
+ diag = tx * step_count
+ first, last = get_merge_begin(
+ source,
+ by * size,
+ aCount,
+ bCount,
+ aStart,
+ bStart,
+ diag,
+ step_count,
+ )
+ # iterate over the output loop
+ serial_merge(
+ source,
+ dest,
+ source_idx,
+ dest_idx,
+ by * size,
+ aCount,
+ bCount,
+ aStart,
+ bStart,
+ kStart,
+ diag,
+ step_count,
+ first,
+ last,
+ )
+
+ with ib.if_scope(even):
+ merge(source, dest, source_idx, dest_idx)
+ with ib.else_scope():
+ merge(dest, source, dest_idx, source_idx)
+
+ def mergesort(source, dest, source_idx, dest_idx, size, width, even):
+ # calculate the start, mid, and end points of this section
+ start = ib.allocate("int64", (1,), name="start", scope="local")
+ middle = ib.allocate("int64", (1,), name="middle", scope="local")
+ end = ib.allocate("int64", (1,), name="end", scope="local")
+
Review comment:
Related to your comment in https://github.com/apache/tvm/pull/7607#issuecomment-792835543, do we need explicit allocations for these variables?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi edited a comment on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-796070741
@mbrookhart Also on the failure from `test_any.py`: The last test case tests on 0 size input, and it fails both rocm and vulkan. On VK I get segfault, on rocm I get `Floating point exception (core dumped)`. Maybe you want to check if your code is safe for empty input?
On the old sort IR, rocm works fine with empty input. VK still gives segfault, that's another story :)
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-796050338
yeah the same fixed-size threading applied to rocm works fine, so this is definitely a spirv issue. I'll see what's going on.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r590653463
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -136,93 +236,223 @@ def compare(a, b):
out = b <= a
return out
- def bottom_up_merge(source, dest, source_idx, dest_idx, start, middle, end, even):
- """
- Merge the two sections of the array assigned to this thread
- """
- # pylint: disable=arguments-out-of-order
- # initialize iterators
- i[0] = start
- j[0] = middle
- # set up indexes
- base_idx = by * size * axis_mul_after + bz
- # iterate over the output loop
- with ib.for_range(0, end - start) as k:
- i_idx = base_idx + i[0] * axis_mul_after
- j_idx = base_idx + j[0] * axis_mul_after
- k_idx = base_idx + (k + start) * axis_mul_after
-
- def swap_values(source, dest, source_idx, dest_idx):
- def assign_i():
- """assign i value to current output"""
- dest[k_idx] = source[i_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[i_idx]
- i[0] += 1
-
- def assign_j():
- """assign j value to current output"""
- dest[k_idx] = source[j_idx]
- if values is not None:
- dest_idx[k_idx] = source_idx[j_idx]
- j[0] += 1
-
- ## if both of the iterators are in range
- with ib.if_scope(tvm.tir.all(i[0] < middle, j[0] < end)):
- # compare them and insert whichever is next into the output
- with ib.if_scope(compare(source[i_idx], source[j_idx])):
- assign_i()
- with ib.else_scope():
- assign_j()
- # otherwise, simply copy the remainder of the valid iterator to the output
- with ib.else_scope():
- with ib.if_scope(i[0] < middle):
- assign_i()
- with ib.else_scope():
- assign_j()
+ # Sort the lower levels of the merge using odd-even sort, it's fast for small inputs
+ lower_lim = tvm.tir.generic.cast(
+ tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(block_size, "float64"))), "int64"
+ )
- # Switch which input is the source and which is the destination each iteration
- with ib.if_scope(even):
- swap_values(source, dest, source_idx, dest_idx)
- with ib.else_scope():
- swap_values(dest, source, dest_idx, source_idx)
-
- def mergesort(source, dest, source_idx, dest_idx, size, width, even):
- # calculate the start, mid, and end points of this section
- start[0] = width * tid
- with ib.if_scope(start[0] < size):
- middle[0] = tvm.te.min(start[0] + tvm.tir.indexdiv(width, 2), size)
- end[0] = tvm.te.min(start[0] + width, size)
- ## merge the start->middle and middle->end arrays
- bottom_up_merge(source, dest, source_idx, dest_idx, start[0], middle[0], end[0], even)
-
- lim = tvm.tir.generic.cast(
+ _odd_even_sort(
+ ib,
+ size,
+ axis_mul_before * axis_mul_after,
+ 1,
+ is_ascend,
+ keys,
+ keys_swap,
+ values,
+ values_swap,
+ )
+
+ upper_lim = tvm.tir.generic.cast(
tvm.tir.ceil(tvm.tir.log2(tvm.tir.generic.cast(size, "float64"))), "int64"
)
- with ib.for_range(0, lim, dtype="int64") as l2_width:
- width = 2 << l2_width
+
+ def get_merge_begin(source, base_idx, aCount, bCount, aStart, bStart, diag, step_count):
+ first = ib.allocate("int64", (1,), name="first", scope="local")
+ mid = ib.allocate("int64", (1,), name="mid", scope="local")
+ last = ib.allocate("int64", (1,), name="last", scope="local")
+ first[0] = tvm.te.max(0, diag - bCount)
+ last[0] = tvm.te.min(diag, aCount)
+ with ib.while_loop(first[0] < last[0]):
+ mid[0] = (first[0] + last[0]) >> 1
+ a = source[base_idx + (aStart + mid[0])]
+ b = source[base_idx + (bStart + diag - 1 - mid[0])]
+ with ib.if_scope(compare(a, b)):
+ first[0] = mid[0] + 1
+ with ib.else_scope():
+ last[0] = mid[0]
+ return first, last
+
+ def serial_merge(
+ source,
+ dest,
+ source_idx,
+ dest_idx,
+ base_idx,
+ aCount,
+ bCount,
+ aStart,
+ bStart,
+ kStart,
+ diag,
+ step_count,
+ first,
+ last,
+ ):
+ i = ib.allocate("int64", (1,), name="i", scope="local")
+ j = ib.allocate("int64", (1,), name="j", scope="local")
+ i[0] = aStart + first[0]
+ j[0] = bStart + diag - last[0]
+ with ib.for_range(0, tvm.te.min(aCount + bCount - diag, step_count)) as count:
+ i_idx = base_idx + i[0]
+ j_idx = base_idx + j[0]
+ k_idx = base_idx + (kStart + diag + count)
+
+ def assign_i():
+ """assign i value to current output"""
+ dest[k_idx] = source[i_idx]
+ if values is not None:
+ dest_idx[k_idx] = source_idx[i_idx]
+ i[0] += 1
+
+ def assign_j():
+ """assign j value to current output"""
+ dest[k_idx] = source[j_idx]
+ if values is not None:
+ dest_idx[k_idx] = source_idx[j_idx]
+ j[0] += 1
+
+ ## if both of the iterators are in range
+ with ib.if_scope(tvm.tir.all(i[0] < aStart + aCount, j[0] < bStart + bCount)):
+ # compare them and insert whichever is next into the output
+ with ib.if_scope(compare(source[i_idx], source[j_idx])):
+ assign_i()
+ with ib.else_scope():
+ assign_j()
+ # otherwise, simply copy the remainder of the valid iterator to the output
+ with ib.else_scope():
+ with ib.if_scope(i[0] < aStart + aCount):
+ assign_i()
+ with ib.else_scope():
+ assign_j()
+
+ with ib.for_range(0, upper_lim - lower_lim, dtype="int64") as l2_width:
+ width = 2 << (l2_width + lower_lim)
# Define and launch the cuda kernel
with ib.new_scope():
- i = ib.allocate("int64", (1,), name="i", scope="local")
- j = ib.allocate("int64", (1,), name="j", scope="local")
- start = ib.allocate("int64", (1,), name="start", scope="local")
- middle = ib.allocate("int64", (1,), name="middle", scope="local")
- end = ib.allocate("int64", (1,), name="end", scope="local")
- tx = te.thread_axis("threadIdx.x")
- bx = te.thread_axis("blockIdx.x")
- ib.scope_attr(tx, "thread_extent", nthread_tx)
- # Reduce the number of blocks as the work per thread grows
- ib.scope_attr(
- bx,
- "thread_extent",
- tvm.tir.generic.cast(ceil_div(size, width * max_threads), "int32"),
- )
- tid = bx * nthread_tx + tx
Review comment:
The slow down is on cuda? We only need this workaround for VK.
If 15 % slow down is a rough estimate compared to an ideal world where spirv supported dynamic thread group size, that doesn't sound too bad to me.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] mbrookhart commented on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
mbrookhart commented on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-795866540
@masahi I pushed the change to vulkan threading, could you test that on your GPU?
I build the CI docker on my machine and attempting to reproduce the dynamic topk failure, I still was unable to. I'm not sure where to go next.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r591921822
##########
File path: python/tvm/topi/cuda/sort.py
##########
@@ -121,41 +221,85 @@ def _sort_common(
## On the final iteration of the algorithm, one thread will merge two sorted lists
## to sort the entire array
- max_threads = int(tvm.target.Target.current(allow_none=False).max_num_threads)
- nthread_tx = max_threads
- nthread_bx = ceil_div(size, max_threads)
- nthread_by = axis_mul_before
- nthread_bz = axis_mul_after
-
- def compare(a, b):
- """
- Compare a and b in proper ascending or descending order
- """
- if is_ascend:
- out = a <= b
- else:
- out = b <= a
- return out
-
- def bottom_up_merge(source, dest, source_idx, dest_idx, start, middle, end, even):
- """
- Merge the two sections of the array assigned to this thread
- """
- # pylint: disable=arguments-out-of-order
- # initialize iterators
- i = ib.allocate("int64", (1,), name="i", scope="local")
- j = ib.allocate("int64", (1,), name="j", scope="local")
- i[0] = start
- j[0] = middle
- # set up indexes
- base_idx = by * size * axis_mul_after + bz
- # iterate over the output loop
- with ib.for_range(0, end - start) as k:
- i_idx = base_idx + i[0] * axis_mul_after
- j_idx = base_idx + j[0] * axis_mul_after
- k_idx = base_idx + (k + start) * axis_mul_after
+ with ib.if_scope(size > 0):
Review comment:
maybe better to move this outside of `_sort_common`? To make indent too deep.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-793081245
@mbrookhart This doesn't work on VK/SPIR-V because they SPIR-V requires a thread block size to be a compile time constant, see https://github.com/apache/tvm/blob/d7f57532746680732e58ab028d8c3129b9140d3d/src/target/spirv/codegen_spirv.cc#L106
In particular, the following size selection has an issue, because `width` is not a constant:
https://github.com/apache/tvm/blob/d2940ca8e25aa400825b79ac7d03da0ff4543db2/python/tvm/topi/cuda/sort.py#L335-L337
So to workaround this problem, can we either
* Use old TIR sort for VK
* (Better) Use a constant thread block size for VK
?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#issuecomment-799954582
I confirmed that this sort also works with VK / SPIRV backend with the following great results:
```
main_vulkan mergepath_vk mergepath_rocm
[1000, 1, 1] 3.04 ms 0.90 ms 0.12 ms
[1, 1000, 1] 3.04 ms 0.88 ms 0.12 ms
[1, 1, 1000] 3.04 ms 0.73 ms 0.11 ms
[1000, 10, 10] 3.49 ms 0.91 ms 0.24 ms
[10, 1000, 10] 3.49 ms 0.91 ms 0.23 ms
[10, 10, 1000] 3.25 ms 0.73 ms 0.22 ms
[1000, 100, 100] 124.09 ms 14.34 ms 19.13 ms
[100, 1000, 100] 186.12 ms 12.03 ms 17.07 ms
[100, 100, 1000] 119.85 ms 10.88 ms 15.87 ms
[10000, 1, 1] 31.27 ms 1.21 ms 0.24 ms
[1, 10000, 1] 31.23 ms 1.18 ms 0.24 ms
[1, 1, 10000] 31.20 ms 1.03 ms 0.23 ms
[10000, 10, 10] 40.54 ms 3.53 ms 2.72 ms
[10, 10000, 10] 41.29 ms 3.43 ms 2.60 ms
[10, 10, 10000] 38.88 ms 3.28 ms 2.52 ms
[10000, 100, 100] 3006.17 ms 270.40 ms 340.12 ms
[100, 10000, 100] 3327.44 ms 207.98 ms 286.45 ms
[100, 100, 10000] 1949.95 ms 179.33 ms 252.66 ms
[100000, 1, 1] 259.56 ms 1.37 ms 0.49 ms
[1, 100000, 1] 259.52 ms 1.30 ms 0.50 ms
[1, 1, 100000] 259.65 ms 1.14 ms 0.49 ms
[100000, 10, 10] 538.90 ms 33.35 ms 34.05 ms
[10, 100000, 10] 514.35 ms 31.01 ms 31.86 ms
[10, 10, 100000] 396.63 ms 29.90 ms 30.78 ms
[1000000, 1, 1] 2336.85 ms 5.68 ms 3.87 ms
[1, 1000000, 1] 2346.90 ms 5.69 ms 3.87 ms
[1, 1, 1000000] 2337.73 ms 5.51 ms 3.83 ms
[1000000, 10, 10] 6797.32 ms 411.09 ms 451.99 ms
[10, 1000000, 10] 6087.86 ms 347.81 ms 387.86 ms
[10, 10, 1000000] 4416.72 ms 333.54 ms 372.60 ms
[(4507,), 0] 15.22 ms 0.87 ms 0.20 ms
[(1, 122640), 1] 290.39 ms 1.18 ms 0.51 ms
[(1, 120000), 1] 285.99 ms 1.14 ms 0.51 ms
[(1, 30000), 1] 72.42 ms 0.98 ms 0.26 ms
[(1, 7500), 1] 18.76 ms 0.83 ms 0.19 ms
[(1, 1000), 1] 3.00 ms 0.68 ms 0.11 ms
```
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a change in pull request #7611: [TOPI][GPU] Mergepath sort with odd-even block sort
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #7611:
URL: https://github.com/apache/tvm/pull/7611#discussion_r589828485
##########
File path: src/tir/transforms/storage_access.cc
##########
@@ -132,6 +132,7 @@ void StorageAccessVisitor::VisitStmt_(const AttrStmtNode* op) {
StmtExprVisitor::VisitStmt_(op);
}
env_threads_.pop_back();
+ } else if (op->attr_key == attr::hand_threaded) {
} else {
Review comment:
How about this?
```
else if (op->attr_key != attr::hand_threaded) {
StmtExprVisitor::VisitStmt_(op);
}
```
And better to leave a comment on what's up with `attr::hand_threaded`
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org