MCPcopy Create free account
hub / github.com/apache/tvm / _flat_product_range

Function _flat_product_range

python/tvm/tirx/exec_context.py:313–346  ·  view source on GitHub ↗
(
    major: AxisRange, lane: AxisRange, lo: int, hi: int
)

Source from the content-addressed store, hash-verified

311
312
313def _flat_product_range(
314 major: AxisRange, lane: AxisRange, lo: int, hi: int
315) -> tuple[AxisRange, AxisRange]:
316 active_min = major.offset * 32 + lane.offset
317 active_max = (
318 (major.offset + major.stride * (major.extent - 1)) * 32
319 + lane.offset
320 + lane.stride * (lane.extent - 1)
321 + 1
322 )
323 if lo <= active_min and active_max <= hi:
324 return major, lane
325
326 if major.stride != 1 or lane.stride != 1:
327 raise ExecContextError("flat thread range narrowing requires unit-stride axes")
328
329 lane_hi = lane.offset + lane.extent
330 major_hi = major.offset + major.extent
331 hit_lo = max(major.offset, (lo - lane_hi) // 32 + 1)
332 hit_hi = min(major_hi, _ceildiv(hi - lane.offset, 32))
333 if hit_hi <= hit_lo:
334 raise ExecContextError("flat thread range produces empty active set")
335
336 if hit_hi == hit_lo + 1:
337 new_lane_lo = max(lane.offset, lo - hit_lo * 32)
338 new_lane_hi = min(lane_hi, hi - hit_lo * 32)
339 if new_lane_hi <= new_lane_lo:
340 raise ExecContextError("flat thread range produces empty lane range")
341 return AxisRange(1, hit_lo), AxisRange(new_lane_hi - new_lane_lo, new_lane_lo)
342
343 if lo <= hit_lo * 32 + lane.offset and (hit_hi - 1) * 32 + lane_hi <= hi:
344 return AxisRange(hit_hi - hit_lo, hit_lo), lane
345
346 raise ExecContextError("flat thread range would require a non-rectangular lane/warp active set")
347
348
349def scope_switch(A: ActiveSet, scope_kind: str) -> Split:

Callers 1

filter_narrowFunction · 0.85

Calls 5

ExecContextErrorClass · 0.85
_ceildivFunction · 0.85
AxisRangeClass · 0.85
maxFunction · 0.50
minFunction · 0.50

Tested by

no test coverage detected

Used in the wild real call sites across dependent graphs

searching dependent graphs…