| 236 | upper_lim = ceil_log2(size) |
| 237 | |
| 238 | def get_merge_begin(source, base_idx, aCount, bCount, aStart, bStart, diag, first, last): |
| 239 | max_val = tvm.te.max(0, diag - bCount) |
| 240 | min_val = tvm.te.min(diag, aCount) |
| 241 | if is_webgpu: |
| 242 | first[0] = cast(max_val, target_dtype) |
| 243 | last[0] = cast(min_val, target_dtype) |
| 244 | else: |
| 245 | first[0] = max_val |
| 246 | last[0] = min_val |
| 247 | with T.While(first[0] < last[0]): |
| 248 | mid = (first[0] + last[0]) >> 1 |
| 249 | a = source[base_idx + (aStart + mid)] |
| 250 | b = source[base_idx + (bStart + diag - 1 - mid)] |
| 251 | with T.If(compare(a, b)): |
| 252 | with T.Then(): |
| 253 | first[0] = mid + 1 |
| 254 | with T.Else(): |
| 255 | last[0] = mid |
| 256 | |
| 257 | def serial_merge( |
| 258 | source, |