apache / tvm

Open deep learning compiler stack for cpu, gpu and specialized accelerators
https://tvm.apache.org/
Apache License 2.0
11.67k stars 3.45k forks source link

[TIR][NarrowDataType] Bufferload's index should not inherit bits constraint of value #17411

Closed wrongtest-intellif closed 1 week ago

wrongtest-intellif commented 1 week ago

Fix an issue in NarrowDataType pass. When we perform narrowing on buffer indices, the integer bits constraint is irrelevant to buffer load value's integer bits. Thus it looks like we should reset the context bit info when traverse into the indices field.

Another small change is in cast node rewrite, skip node construction when dtype actually matches.

tqchen commented 1 week ago

do u mind leave a quick comment about what will happen without this PR

wrongtest-intellif commented 1 week ago

do u mind leave a quick comment about what will happen without this PR

In the main branch, this case's indices would not change after narrowing with i32.

@T.prim_func
def before(A: T.Buffer((16,), "int64")):
    for i in range(T.int64(15)):
          A[i + T.int64(1)] = A[i] + T.int64(1)

The deduction chain is like below:

  1. A[i] + T.int64(1) -> 64b
  2. A[i] -> 64b
  3. indices of load A -> 64b (here is the issue)
  4. i -> requires 64b