Support dynamic indexing in ti.Vector/ti.Matrix

[k-ye]

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

[strongoier]

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:

• Enforce each sub-field has the same shape. ([Lang] Enforce members of a matrix field to have same shape #3761)
• Develop an algorithm to identify the cases where the dynamic indexing can be turned into the form of pointer + offset at compile time. ([RFC] Algorithm for identifying the cases where dynamic indexing can be turned into pointer + offset #3810)
• Enable dynamic indexing when the above algorithm shows feasibility, and report an error otherwise. ([Lang] Enable dynamic indexing of matrix field elements when possible #3865)

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]

Close as this has been finished.