Skip to content

[Tracking Issue] Inherit block iter_var from parent blocks #71

@yzh119

Description

@yzh119

Pitch

Sparse Iteration Lowering pass would not create inherit block iter_vars from outer blocks. This is okay but does not follow the design principle of TensorIR that blocks should be pluggable (if we do not inherit outer block iter_vars, then this block relies on its context and cannot be compute_at to an arbitrary block).

# lowered nested block structure in previous design, we can use `vi` inside "inner" block, however, we cannot `compute_at` "inner" outside "outer" because information regarding "vi" is lost.
for i in range(20):
    with block("outer"):
        vi = T.axis.spatial(i, 20)
        for j in range(indptr[i + 1] - indptr[i]):
        with block("inner"):
            vj = T.axis.spatial(j, 10)
            ...

# inherit `vi` from "outer" block in inner block, we can `compute_at` "inner" block outside "outer" and `vi_1` would be re-bound to other iter values.
for i in range(20):
    with block("outer"):
        vi = T.axis.spatial(i, 20)
        for j in range(indptr[i + 1] - indptr[i]):
        with block("inner"):
            vi_1 = T.axis.spatial(vi, 20)
            vj = T.axis.spatial(j, 10)
            ...

This design also results in bugs such as #60 : if we carry all outer block iter_vars, we do not need to consider iter_vars in outer blocks in cache read/write.

Milestone

  • Fix the behavior of sparse iteration lowering to inherit iter_var from outer blocks.
  • Checks whether examples would break or not.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type
    No fields configured for issues without a type.

    Projects

    Status

    TODO

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions