[RFC] Taichi Matrix & Vector refactor plan

[jim19930609]

Background and Purpose

Currently, Taichi scalarizes all the Matrix and corresponding operations at Python level. This approach has major drawbacks as follow:

1. No actual SIMD support for Matrix data type during CodeGen.
2. Lack of Matrix information at CHI IR level, thus not able to write any optimization pass wrt Matrix operations.
3. Huge overhead to migrate Matrix support to another Frontend Language (Julia, Javascript, etc)

To address the above mentioned drawbacks, we describe our refactor goal as follow:

1. Remove early scalarization at Python-level and make full Matrix information visible to CHI IR level.
2. Enable the ability to optimize Matrix operations via CHI IR Passes.
3. Enable SIMD instructions code generation for Matrix operations.
4. Minimize Matrix handling at Python-level to enhance its multiple Frontend Language support
5. No regression wrt functionality, performance and compilation efficiency especially with dynamic indexing on

Design Overview

In general there are three separate code paths with Matrix involved:

1. local Matrix
2. MatrixField：Field + Matrix
3. MatrixNdarray：Ndarray + Matrix

Each of the code path involves four typical use cases:

1. Indexed Matrix operations：
   • Matrix Indexing
   • Indexed Matrix Access：Load, Store, Add, Mul, ...
2. MatrixType operations:
   • Matrix initialization
   • Matrix operations：Load, Store, matmul, transpose, atomic, ...

Design Details

Local Matrix

Implementation plan: #5478

1. Representation

ti.Matrix returns a Expr(MatrixExpression) with TensorType

@ti.kernel
def example():
    x = ti.Matrix([[1, 2], [3, 4]], dt=ti.f32)
    print(x)

2. IR changes

[New Frontend IR] MatrixExpression: Initializing a matrix with ti.Matrix will result in the following Frontend IR:

kernel {
  $0 = alloca @tmp0
  @tmp0 = [(cast_value<f32> 1.0), (cast_value<f32> 2.0), (cast_value<f32> 3.0), (cast_value<f32> 4.0)] (dt=[Tensor (2, 2) f32])
  $2 = alloca @tmp1
  @tmp1 = @tmp0
  print @tmp1, "\n"
}

[New CHI IR] MatrixInitStmt: Matrix initialization
[Reuse CHI IR] PtrOffsetStmt: Represent Matrix indexing

3. Codegen (LLVM)

MatrixInitStmt -> llvm::VectorType（Init + InsertElement）
BinaryOpStmt -> Same LLVM Binary Instruction with llvm::VectorType operands

4. Lowering details

MatrixField

AOS Memory Layout

1. AOS Python Interface

Current：ti.Matrix.field(m=2, n=2, shape=3, ti.AOS)
After SNodeTreeType：builder.add_field(dtype=ti.vec3, name='x', shape=(4, 8))

2. Representation

• The leaf snode of a SNodeTree becomes Place SNode with MatrixType, which get compiled into llvm::VectorType during CodeGen
• Operations to MatrixField get divided into two stages:
  • Obtain Place SNode through traversing SNodeTree
  • Apply operations to the MatrixTyped data stored in the Place SNode.

3. Lowering details

SOA/Customized Memory Layout

1. SOA Python Interface

Current：

x = Matrix.field(m=2, n=2, shape=3, ti.SOA)
    or
x = Matrix.field(m=2, n=2)
ti.root.....place(x.get_scalar_field(0, 0))
ti.root.....place(x.get_scalar_field(0, 1))
ti.root.....place(x.get_scalar_field(1, 0))
ti.root.....place(x.get_scalar_field(1, 1))

After SNodeTreeType:

builder = ti.SNodeTreeBuilder()
builder.add_field(dtype=ti.f32, name='v0')
builder.add_field(dtype=ti.f32, name='v1')
builder.add_field(dtype=ti.f32, name='v2')
for v in ['v0', 'v1', 'v2']:
  builder.tree().dense(ti.ij, 4).place(v)

builder.add_field_view(dtype=ti.vec3, name='vel', components=['v0', 'v1', 'v2'])

2. Representation

• [New Frontend IR] ListOfSNodeExpression(): Contains a list of SNodes, representing SOA-MatrixField at Frontend IR.
• Support MatrixType with stride to represent SOA-Matrix, so SOA-Matrix can be represented the same way as AOS-Matrix at CHi IR level
• SOA-MatrixField will get scalarized during CodeGen.

3. Lowering details

MatrixNdarray

MatrixNdarray share similar implementations wrt MatrixField. We mainly describe the lowering details in this section.

AOS Memory Layout

Lowering details:

SOA Memory Layout

Lowering details:

Fallback & Scalarization

• [LocalMatrix & AOS] Fallback to "Array"：
  • AOS fallback at CodeGen level
  • ArrayType is friendly to dynamic indexing
  • Possible reasons to fallback：Lack SIMD support at hardware level; Lack SIMD for specific dtype or size; performance regression with SIMD
• [SOA] Fallback to standalone scalars：
  • SOA get scalarized during CodeGen

Roadmap

Proposed Demos

1. python/taichi/examples/simulation/mpm88.py
2. python/taichi/examples/simulation/nbody.py