Initialize the output buffers by copying from inputs
(shape, axis, keys_in, keys_out, values_out=None, value_init_func=None)
| 37 | |
| 38 | |
| 39 | def _sort_init(shape, axis, keys_in, keys_out, values_out=None, value_init_func=None): |
| 40 | """Initialize the output buffers by copying from inputs""" |
| 41 | axis_mul_before = 1 |
| 42 | axis_mul_after = 1 |
| 43 | if axis < 0: |
| 44 | axis = len(shape) + axis |
| 45 | for i, value in enumerate(shape, 0): |
| 46 | if i < axis: |
| 47 | axis_mul_before *= value |
| 48 | elif i > axis: |
| 49 | axis_mul_after *= value |
| 50 | |
| 51 | # Set up threading |
| 52 | max_threads = int(tvm.target.Target.current(allow_none=False).attrs["max_num_threads"]) |
| 53 | nthread_tx = max_threads |
| 54 | nthread_bx = ceil_div(shape[axis], max_threads) |
| 55 | nthread_by = axis_mul_before * axis_mul_after |
| 56 | |
| 57 | # Copy the keys_in to initial output |
| 58 | tx, bx, by, ntx, nbx, nby = _get_threads(nthread_tx, nthread_bx, nthread_by) |
| 59 | with T.frame_scope( |
| 60 | [ |
| 61 | T.attr(tx, "thread_extent", ntx), |
| 62 | T.attr(bx, "thread_extent", nbx), |
| 63 | T.attr(by, "thread_extent", nby), |
| 64 | ] |
| 65 | ): |
| 66 | tid = bx * nthread_tx + tx |
| 67 | by_val = by % axis_mul_before |
| 68 | bz = by // axis_mul_before |
| 69 | idx = (by_val * shape[axis] + tid) * axis_mul_after + bz |
| 70 | with T.If(tid < shape[axis]): |
| 71 | with T.Then(): |
| 72 | keys_out[idx] = keys_in[idx] |
| 73 | if values_out is not None: |
| 74 | values_out[idx] = value_init_func(idx, tid) |
| 75 | |
| 76 | return axis_mul_before, axis_mul_after |
| 77 | |
| 78 | |
| 79 | ## TODO(mbrookhart): These are effective optimziation hyperparametrs |
no test coverage detected
searching dependent graphs…