Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support dynamic indexing in ti.Vector/ti.Matrix #2590

Closed
k-ye opened this issue Jul 27, 2021 · 2 comments
Closed

Support dynamic indexing in ti.Vector/ti.Matrix #2590

k-ye opened this issue Jul 27, 2021 · 2 comments
Assignees
Labels
feature request Suggest an idea on this project

Comments

@k-ye
Copy link
Member

k-ye commented Jul 27, 2021

The original issue was #1004. This issue tracks the progress of supporting the new feature. Detailed design stories are commented below.

@strongoier
Copy link
Contributor

strongoier commented Dec 8, 2021

Dynamic indexing refers to the following two cases:

@ti.kernel
def dynamic_index_local_vector(i: ti.i32) -> ti.f32:
    a = ti.Vector([0.0, 1.0, 2.0])
    return a[i]

b = ti.Vector.field(3, ti.f32, shape=())

@ti.kernel
def dynamic_index_vector_field(i: ti.i32) -> ti.f32:
    return b[None][i]

We would like to gradually support this feature with ti.init(dynamic_index=True).

  • If we can finally confirm that this feature will not harm performance of non-dynamic-indexing code, we will make this behavior as default.

A few months ago, @squarefk made a prototype for this in #2543, #2637, #2673, #2773. He also gave a talk (https://github.com/taichi-dev/taichicon/releases/download/taichicon02/Dynamic.Indexing-Freely.Access.Vector.and.Matrix.Elements.at.Runtime-Yu.Fang.pdf) at TaichiCon02 to present his ideas - common use cases can be translated to the form of pointer + offset by the compiler. Here are some important technical details of his implementation:

  • TensorType is added in CHI IR for vectors and matrices;
  • PtrOffsetStmt is added in CHI IR for accessing a memory address of form pointer + offset;
  • TensorElementExpression is added in frontend IR to represent a vector/matrix with dynamic indices.

However, there are still lots of engineering efforts to be done in order to make this idea fully fit into the language and the compiler. Recently, I've been driving the development of this feature. Related work includes #3218, #3228, #3237, #3244, #3468, #3517, #3524, #3546, #3599. Also, this feature pushes us to support frontend type check (#3301). Now we can completely use dynamic indexing of local vectors/matrices now.

My next focus is to fully support dynamic indexing of elements of vector/matrix fields. A n-D vector field is currently implemented as a tuple of n scalar fields, and we will refer to them as sub-fields below. Here is the preliminary roadmap:

Handle other rare use cases with a fallback plan:

# assume v is an element of a vector field
for i in ti.static(range(v.n)):
    if i == index:
        return v[i]

strongoier added a commit that referenced this issue Dec 27, 2022
Issue: #2590

### Brief Summary

1. Dynamic indexing in metal is supported by mimicking the
implementation of that in LLVM backends.
2. `ret_type` change of `ArgLoadStmt` is moved out of the
`ScalarizePointers` pass, which is designed to scalarize `alloca` and
will be removed once dynamic indexing becomes default.
3. Configs of dynamic indexing tests are fixed.

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
strongoier added a commit that referenced this issue Dec 30, 2022
Issue: #2590

### Brief Summary

This PR supports dynamic indexing in spirv (sister PR: #6985).
`_test_local_matrix_non_constant_index()` is modified due to lack of
`AssertStmt` support in spirv.

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
strongoier added a commit that referenced this issue Jan 5, 2023
Issue: #2590

### Brief Summary

Under pure `dynamic_index` setting, `MatrixPtrStmt`s are not scalarized.
It actually produces `2n` more instructions (`n` `ConstStmt`s and n
`MatrixPtrStmt`s) than the scalarized setting, where `n` is the number
of usages of `MatrixPtrStmt`s. This PR adds `ExtractPointers` pass to
eliminate all the redundant instructions. See comments in the code for
details.

After this PR, the number of instructions after the `scalarize()` pass
of the script in #6933 under dynamic index reduces from 49589 to 26581,
and the compilation time reduces from 20.02s to 7.82s.

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
feisuzhu pushed a commit to feisuzhu/taichi that referenced this issue Jan 5, 2023
Issue: taichi-dev#2590

### Brief Summary

This PR supports dynamic indexing in spirv (sister PR: taichi-dev#6985).
`_test_local_matrix_non_constant_index()` is modified due to lack of
`AssertStmt` support in spirv.

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
feisuzhu pushed a commit to feisuzhu/taichi that referenced this issue Jan 5, 2023
Issue: taichi-dev#2590

### Brief Summary

Under pure `dynamic_index` setting, `MatrixPtrStmt`s are not scalarized.
It actually produces `2n` more instructions (`n` `ConstStmt`s and n
`MatrixPtrStmt`s) than the scalarized setting, where `n` is the number
of usages of `MatrixPtrStmt`s. This PR adds `ExtractPointers` pass to
eliminate all the redundant instructions. See comments in the code for
details.

After this PR, the number of instructions after the `scalarize()` pass
of the script in taichi-dev#6933 under dynamic index reduces from 49589 to 26581,
and the compilation time reduces from 20.02s to 7.82s.

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
lin-hitonami added a commit that referenced this issue Jan 5, 2023
Issue: #2590

### Brief Summary

Previously `ScalarizePointers` (which is for allocas) takes place only
when `dynamic_index=False`. In this PR, we automatically identify those
allocas only indexed with constants and get them scalarized, and let the
remaining allocas go through the code path for dynamic indexing. In this
way, we make old user programs behave the same regardless of the
`dynamic_index` option, and provide dynamic indexing support for new
user programs at the same time.

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lin Jiang <[email protected]>
lin-hitonami pushed a commit that referenced this issue Jan 5, 2023
Issue: #2590

### Brief Summary

Dynamic indexing is now automatically supported and no switch is needed.

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
PGZXB pushed a commit to PGZXB/taichi that referenced this issue Jan 6, 2023
Issue: taichi-dev#2590

### Brief Summary

Dynamic indexing is now automatically supported and no switch is needed.

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
lin-hitonami pushed a commit that referenced this issue Jan 17, 2023
Issue: #2590

### Brief Summary

The `dynamic_index` switch of `ti.init()` should be removed according to
the [deprecation
notice](https://github.com/taichi-dev/taichi/releases/tag/v1.4.0).
@strongoier
Copy link
Contributor

Close as this has been finished.

quadpixels pushed a commit to quadpixels/taichi that referenced this issue May 13, 2023
Issue: taichi-dev#2590

### Brief Summary

1. Dynamic indexing in metal is supported by mimicking the
implementation of that in LLVM backends.
2. `ret_type` change of `ArgLoadStmt` is moved out of the
`ScalarizePointers` pass, which is designed to scalarize `alloca` and
will be removed once dynamic indexing becomes default.
3. Configs of dynamic indexing tests are fixed.

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
quadpixels pushed a commit to quadpixels/taichi that referenced this issue May 13, 2023
Issue: taichi-dev#2590

### Brief Summary

This PR supports dynamic indexing in spirv (sister PR: taichi-dev#6985).
`_test_local_matrix_non_constant_index()` is modified due to lack of
`AssertStmt` support in spirv.

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
quadpixels pushed a commit to quadpixels/taichi that referenced this issue May 13, 2023
Issue: taichi-dev#2590

### Brief Summary

Under pure `dynamic_index` setting, `MatrixPtrStmt`s are not scalarized.
It actually produces `2n` more instructions (`n` `ConstStmt`s and n
`MatrixPtrStmt`s) than the scalarized setting, where `n` is the number
of usages of `MatrixPtrStmt`s. This PR adds `ExtractPointers` pass to
eliminate all the redundant instructions. See comments in the code for
details.

After this PR, the number of instructions after the `scalarize()` pass
of the script in taichi-dev#6933 under dynamic index reduces from 49589 to 26581,
and the compilation time reduces from 20.02s to 7.82s.

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
quadpixels pushed a commit to quadpixels/taichi that referenced this issue May 13, 2023
Issue: taichi-dev#2590

### Brief Summary

Previously `ScalarizePointers` (which is for allocas) takes place only
when `dynamic_index=False`. In this PR, we automatically identify those
allocas only indexed with constants and get them scalarized, and let the
remaining allocas go through the code path for dynamic indexing. In this
way, we make old user programs behave the same regardless of the
`dynamic_index` option, and provide dynamic indexing support for new
user programs at the same time.

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lin Jiang <[email protected]>
quadpixels pushed a commit to quadpixels/taichi that referenced this issue May 13, 2023
Issue: taichi-dev#2590

### Brief Summary

Dynamic indexing is now automatically supported and no switch is needed.

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
quadpixels pushed a commit to quadpixels/taichi that referenced this issue May 13, 2023
quadpixels pushed a commit to quadpixels/taichi that referenced this issue May 13, 2023
Issue: taichi-dev#2590

### Brief Summary

The `dynamic_index` switch of `ti.init()` should be removed according to
the [deprecation
notice](https://github.com/taichi-dev/taichi/releases/tag/v1.4.0).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
feature request Suggest an idea on this project
Projects
None yet
Development

No branches or pull requests

3 participants