How to use the taichi.i32 function in taichi

def test_pass_by_value():

  @ti.func
  def set_val(x, i):
    x = i

  ret = ti.var(ti.i32, shape=())

  @ti.kernel
  def task():
    set_val(ret[None], 112)

  task()
  assert ret[None] == 0

def test_parallel_range_for():
  n = 1024 * 1024
  val = ti.var(ti.i32, shape=(n))

  @ti.kernel
  def fill():
    ti.parallelize(8)
    ti.block_dim(8)
    for i in range(n):
      val[i] = i

  fill()

  for i in range(n):
    assert val[i] == i

def test_pointer2():
  x = ti.var(ti.f32)
  s = ti.var(ti.i32)

  n = 128

  @ti.layout
  def place():
    ti.root.dense(ti.i, n).pointer().dense(ti.i, n).place(x)
    ti.root.place(s)

  @ti.kernel
  def activate():
    for i in range(n * n):
      x[i] = i

  @ti.kernel
  def func():
    for i in x:

def test_kernel_template_mapper():
  x = ti.var(ti.i32)
  y = ti.var(ti.f32)

  @ti.layout
  def layout():
    ti.root.place(x, y)

  mapper = ti.KernelTemplateMapper(
      (ti.template(), ti.template(), ti.template()),
      template_slot_locations=(0, 1, 2))
  assert mapper.lookup((0, 0, 0)) == 0
  assert mapper.lookup((0, 1, 0)) == 1
  assert mapper.lookup((0, 0, 0)) == 0
  assert mapper.lookup((0, 0, 1)) == 2
  assert mapper.lookup((0, 1, 0)) == 1

  mapper = ti.KernelTemplateMapper((ti.i32, ti.i32, ti.i32), ())

import os
import numpy as np
import cv2
import math
import time
import random
from renderer_utils import out_dir, ray_aabb_intersection, inf, eps, \
  intersect_sphere, sphere_aabb_intersect_motion, inside_taichi
import sys

res = 1280, 720
num_spheres = 1024
color_buffer = ti.Vector(3, dt=ti.f32)
sphere_pos = ti.Vector(3, dt=ti.f32)
bbox = ti.Vector(3, dt=ti.f32)
grid_density = ti.var(dt=ti.i32)
voxel_has_particle = ti.var(dt=ti.i32)
max_ray_depth = 4
use_directional_light = True

particle_x = ti.Vector(3, dt=ti.f32)
particle_v = ti.Vector(3, dt=ti.f32)
particle_color = ti.var(ti.i32)
pid = ti.var(ti.i32)
num_particles = ti.var(ti.i32)

fov = 0.23
dist_limit = 100

exposure = 1.5
camera_pos = ti.Vector([0.5, 0.32, 1.7])
vignette_strength = 0.9

dir /= length

        rotated_x = dir[0] * ti.cos(angle) + dir[2] * ti.sin(angle)
        rotated_z = -dir[0] * ti.sin(angle) + dir[2] * ti.cos(angle)
        dir[0] = rotated_x
        dir[2] = rotated_z
        point = camera_origin + (k + 1) * dx * dir

        # Convert to coordinates of the density grid box
        box_x = point[0] + 0.5
        box_y = point[1] + 0.5
        box_z = point[2] + 0.5

        # Density grid location
        index_x = ti.cast(ti.floor(box_x * density_res), ti.i32)
        index_y = ti.cast(ti.floor(box_y * density_res), ti.i32)
        index_z = ti.cast(ti.floor(box_z * density_res), ti.i32)
        index_x = ti.max(0, ti.min(index_x, density_res - 1))
        index_y = ti.max(0, ti.min(index_y, density_res - 1))
        index_z = ti.max(0, ti.min(index_z, density_res - 1))

        flag = 0
        if in_box(point[0], point[1], point[2]):
          flag = 1

        contribution = density[index_z, index_y, index_x] * flag

        ti.atomic_add(field[view_id, y, x], contribution)

def advect(field: ti.template(), field_out: ti.template(),
           t_offset: ti.template(), t: ti.i32):
  """Move field smoke according to x and y velocities (vx and vy)
     using an implicit Euler integrator."""
  for y in range(n_grid):
    for x in range(n_grid):
      center_x = y - v[t + t_offset, y, x][0]
      center_y = x - v[t + t_offset, y, x][1]

      # Compute indices of source cell
      left_ix = ti.cast(ti.floor(center_x), ti.i32)
      top_ix = ti.cast(ti.floor(center_y), ti.i32)

      rw = center_x - left_ix  # Relative weight of right-hand cell
      bw = center_y - top_ix  # Relative weight of bottom cell

      # Wrap around edges
      # TODO: implement mod (%) operator
      left_ix = imod(left_ix, n_grid)
      right_ix = left_ix + 1
      right_ix = imod(right_ix, n_grid)
      top_ix = imod(top_ix, n_grid)
      bot_ix = top_ix + 1
      bot_ix = imod(bot_ix, n_grid)

      # Linearly-weighted sum of the 4 surrounding cells
      field_out[t, y, x] = (1 - rw) * (

import taichi as ti
import taichi as tc

tc.set_gdb_trigger(True)

# x, y = ti.var(ti.f32), ti.var(ti.f32)
# z, w = ti.var(ti.f32), ti.var(ti.f32)
val = ti.var(ti.i32)
y = ti.var(ti.i32)

ti.cfg.use_llvm = True
# ti.cfg.print_ir = True
# ti.cfg.print_struct_llvm_ir = True
# ti.cfg.print_kernel_llvm_ir = True

# ti.cfg.print_ir = True
# ti.runtime.print_preprocessed = True

n = 32

@ti.layout
def values():
  # fork = ti.root.dense(ti.k, 128)
  # fork.dense(ti.ij, 16).place(x, y)

assert self.dim in (
            2, 3), "MPM solver supports only 2D and 3D simulations."

        self.res = res
        self.n_particles = ti.var(ti.i32, shape=())
        self.dx = size / res[0]
        self.inv_dx = 1.0 / self.dx
        self.default_dt = 2e-2 * self.dx / size * dt_scale
        self.p_vol = self.dx**self.dim
        self.p_rho = 1000
        self.p_mass = self.p_vol * self.p_rho
        self.max_num_particles = max_num_particles
        self.gravity = ti.Vector(self.dim, dt=ti.f32, shape=())
        self.source_bound = ti.Vector(self.dim, dt=ti.f32, shape=2)
        self.source_velocity = ti.Vector(self.dim, dt=ti.f32, shape=())
        self.pid = ti.var(ti.i32)
        # position
        self.x = ti.Vector(self.dim, dt=ti.f32)
        # velocity
        self.v = ti.Vector(self.dim, dt=ti.f32)
        # affine velocity field
        self.C = ti.Matrix(self.dim, self.dim, dt=ti.f32)
        # deformation gradient
        self.F = ti.Matrix(self.dim, self.dim, dt=ti.f32)
        # material id
        self.material = ti.var(dt=ti.i32)
        self.color = ti.var(dt=ti.i32)
        # plastic deformation volume ratio
        self.Jp = ti.var(dt=ti.f32)

        if self.dim == 2:
            indices = ti.ij

def p2g(f: ti.i32):
  for p in range(0, n_particles):
    base = ti.cast(x[f, p] * inv_dx - 0.5, ti.i32)
    fx = x[f, p] * inv_dx - ti.cast(base, ti.i32)
    w = [0.5 * ti.sqr(1.5 - fx), 0.75 - ti.sqr(fx - 1), 0.5 * ti.sqr(fx - 0.5)]
    new_F = (ti.Matrix.diag(dim=2, val=1) + dt * C[f, p]) @ F[f, p]
    F[f + 1, p] = new_F
    J = ti.determinant(new_F)
    r, s = ti.polar_decompose(new_F)
    cauchy = 2 * mu * (new_F - r) @ ti.transposed(new_F) + \
             ti.Matrix.diag(2, la * (J - 1) * J)
    stress = -(dt * p_vol * 4 * inv_dx * inv_dx) * cauchy
    affine = stress + p_mass * C[f, p]
    for i in ti.static(range(3)):
      for j in ti.static(range(3)):
        offset = ti.Vector([i, j])
        dpos = (ti.cast(ti.Vector([i, j]), real) - fx) * dx
        weight = w[i](0) * w[j](1)
        grid_v_in[base + offset] += weight * (p_mass * v[f, p] + affine @ dpos)
        grid_m_in[base + offset] += weight * p_mass