[IR] Support local tensor

[squarefk]

Related issue = #2590

IR

Python:

a = ti.Vector([1., 2., 3.])

print(a[?])

Initial IR:

<[TensorType> $1 = alloca @tmp     // FrontendAllocaStmt(IdExpression, shape, element_type)
@tmp[0] = elements[0]              // FrontendAssignStmt(LocalTensorElementExpression(IdExpression, indices), elements)
@tmp[1] = elements[1]
@tmp[2] = elements[2]
$? = eval @tmp[?]                  // FrontendEvalStmt
print $?

Lower AST (eliminate IdExpression):

<TensorType> $1 = alloca @tmp      // AllocaStmt
<Type*> $2 = shift ptr [$1 + $?]   // ShiftGlobalPtrStmt
<Type> $3 = local store [$2 <- $?] // LocalStore

<Type> $4 = local load [$2]        // LocalLoad

Offloaded:
AllocaStmt may be converted into GlobalTemporaryStmt or just keep AllocaStmt

Issues

Currently, there are three issues:

• Offload pass only deal with ConstStmt type for offset in PtrOffsetStmt. This will introduce some code restrictions but should be easy to resolve. . Done with another PR.
• During make_local_block pass, Alloca TensorType need LocalStore to initialize which introduces extra difficulty. Hence, it is disabled.
• Loop index keep using individual variable instead of TensorType for now. We may want to change this with another PR. At least the following aspects should be considered.
  2. for I in ti.grouped(ti.ndrange()): need to make sure Alloca TensorType inside loop body.
  3. Loop indices are looked up by Alloca id in some places. Extra work is needed to extract different dim indices from the same Alloca TensorType

[squarefk]

/format

[squarefk]

/format

[squarefk]

/format

[squarefk]

/format

[squarefk]

/format