taichi-dev / taichi

Productive, portable, and performant GPU programming in Python.
https://taichi-lang.org
Apache License 2.0
25.05k stars 2.26k forks source link

Declare ndarray inside kernel from a pointer #8507

Open BernardoCovas opened 2 months ago

BernardoCovas commented 2 months ago

Hi! Is there a way that I can have a list of integers inside a kernel representing pointers to array data of known sizes? Curretly I am working with big image arrays of different sizes that barely fit in memory, and thus I can not copy them into some sort of contiguous array due to copy delays and memory space. Moreover, I will be loading and unloading these images in a dynamic way. I would like to use taichi for processing, but currently I have not found a way to declare inside a taichi kernel that an ineger is an array pointer. Which image I am going to access will depend on computation, and will vary from iteration to iteration inside the same kernel. What I wanted to achive is something like the following:


image0 = ti.ndarray(...) # 10000 x 10000 image
image1 = ti.ndarray(...) # 10000 x 10000 image
image2 = ti.ndarray(...) # 10000 x 10000 image

image_source = ti.ndarray(...) # 10000 x 10000 image

@ti.kernel
def kn(image_src, image_trgts, image_trgts_shape):
   h = image_src.shape[0]
   w = image_src.shape[1]

  for i in range(h):
    for j in range(w):
      target_image_index = (...) # target image index will depend on some computation for the current image pixel
      target_x = (...) # target image coordinates will depend on some other computation
      target_y = (...) # target image coordinates will depend on some other computation

      # get the sape of target image from the shapes array
      tgt_h = image_trgts_shape[target_image, 0]
      tgt_w = image_trgts_shape[target_image, 1]

     # declare the array for the target image that I want to access
     target = ti.ndarray(pointer=image_trgts[target_image], shape=(tgt_h, tgt_w, 3))
     pixel_r = target[target_y, target_x, 0]
     pixel_g = target[target_y, target_x, 1]
     pixel_b = target[target_y, target_x, 2]
     # finally work with target image pixel
     (...)

The real computation has plenty more steps, however this example outlines the issue.

Thank you in advance