taichi-dev / taichi

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

returning a StructType from a kernel hits assert. Is this meant to be possible? #6862

Open oliver-batchelor opened 1 year ago

oliver-batchelor commented 1 year ago

Describe the bug I'm trying to return a StructType from a kernel but running into an assert. Is this meant to be possible? [Taichi] version 1.3.0, llvm 15.0.4, commit 0f25b95e, linux, python 3.10.8

To Reproduce

import taichi as ti
import numpy as np

@ti.dataclass
class AABox:
  """An axis aligned bounding box in 3D space."""
  lower: ti.math.vec3
  upper: ti.math.vec3

@ti.kernel
def point_bounds(points:ti.template()) -> AABox:
  b = AABox(points[0], points[0])
  for i in points:
    b.lower = ti.math.min(b.lower, points[i])
    b.upper = ti.math.max(b.upper, points[i])
  return b
ti.init()

x = ti.Vector.field(3, dtype=ti.f32, shape=10)
points = np.arange(0, 30, dtype=np.float32).reshape(10, 3)
x.from_numpy(points)

b = point_bounds(x)
print(b)

Log/Screenshots

taichi.lang.exception.TaichiCompilationError: 
File "/home/oliver/sync/python_geometry/test.py", line 20, in point_bounds:
  return b
  ^^^^^^^^
Traceback (most recent call last):
  File "/home/oliver/miniconda3/envs/torch13/lib/python3.10/site-packages/taichi/lang/ast/ast_transformer_utils.py", line 26, in __call__
    return method(ctx, node)
  File "/home/oliver/miniconda3/envs/torch13/lib/python3.10/site-packages/taichi/lang/ast/ast_transformer.py", line 759, in build_Return
    expr.make_expr_group(values._members))
  File "/home/oliver/miniconda3/envs/torch13/lib/python3.10/site-packages/taichi/lang/expr.py", line 163, in make_expr_group
    assert i.local_tensor_proxy is not None
AssertionError

Additional comments [Taichi] version 1.3.0, llvm 15.0.4, commit 0f25b95e, linux, python 3.10.8


Taichi Programming Language


Docs: https://docs.taichi-lang.org/ GitHub: https://github.com/taichi-dev/taichi/ Forum: https://forum.taichi.graphics/

Taichi system diagnose:

python: 3.10.8 | packaged by conda-forge | (main, Nov 22 2022, 08:26:04) [GCC 10.4.0] system: linux executable: /home/oliver/miniconda3/envs/torch13/bin/python3.10 platform: Linux-5.4.0-135-generic-x86_64-with-glibc2.31 architecture: 64bit ELF uname: uname_result(system='Linux', node='oliver-laptop', release='5.4.0-135-generic', version='#152-Ubuntu SMP Wed Nov 23 20:19:22 UTC 2022', machine='x86_64') locale: en_NZ.UTF-8 PATH: /home/oliver/miniconda3/envs/torch13/bin:/home/oliver/miniconda3/condabin:/home/oliver/.local/bin:/usr/local/cuda/bin:/home/oliver/.local/bin:/opt/spinnaker/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin PYTHONPATH: ['/home/oliver/miniconda3/envs/torch13/bin', '/home/oliver/miniconda3/envs/torch13/lib/python310.zip', '/home/oliver/miniconda3/envs/torch13/lib/python3.10', '/home/oliver/miniconda3/envs/torch13/lib/python3.10/lib-dynload', '/home/oliver/miniconda3/envs/torch13/lib/python3.10/site-packages', '/home/oliver/sync/py-structs', '/home/oliver/sync/camera_geometry_python', '/home/oliver/sync/open3d_vis', '/home/oliver/sync/python_geometry']

No LSB modules are available. Distributor ID: Ubuntu Description: Ubuntu 20.04.4 LTS Release: 20.04 Codename: focal

import: <module 'taichi' from '/home/oliver/miniconda3/envs/torch13/lib/python3.10/site-packages/taichi/init.py'>

cc: False cpu: True metal: False opengl: True cuda: True MESA-INTEL: warning: Performance support disabled, consider sysctl dev.i915.perf_stream_paranoid=0

vulkan: True

glewinfo not available: [Errno 2] No such file or directory: 'glewinfo'

Sun Dec 11 02:32:02 2022
+-----------------------------------------------------------------------------+ | NVIDIA-SMI 520.61.05 Driver Version: 520.61.05 CUDA Version: 11.8 | |-------------------------------+----------------------+----------------------+ | GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC | | Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. | | | | MIG M. | |===============================+======================+======================| | 0 NVIDIA GeForce ... On | 00000000:01:00.0 Off | N/A | | N/A 55C P8 10W / N/A | 10MiB / 8192MiB | 0% Default | | | | N/A | +-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+ | Processes: | | GPU GI CI PID Type Process name GPU Memory | | ID ID Usage | |=============================================================================| | 0 N/A N/A 1726 G /usr/lib/xorg/Xorg 4MiB | | 0 N/A N/A 2907 G /usr/lib/xorg/Xorg 4MiB | +-----------------------------------------------------------------------------+

[Taichi] version 1.3.0, llvm 15.0.4, commit 0f25b95e, linux, python 3.10.8

[W 12/11/22 02:32:05.871 2903120] [llvm_offline_cache.cpp:load_meta_data@138] Lock /home/oliver/.cache/taichi/ticache/llvm/metadata.lock failed [Taichi] version 1.3.0, llvm 15.0.4, commit 0f25b95e, linux, python 3.10.8 [Taichi] Starting on arch=x64 [W 12/11/22 02:32:06.275 2903120] [offline_cache.h:run@176] Lock /home/oliver/.cache/taichi/ticache/llvm/metadata.lock failed

[Taichi] version 1.3.0, llvm 15.0.4, commit 0f25b95e, linux, python 3.10.8 [Taichi] Starting on arch=opengl

[W 12/11/22 02:32:09.912 2903149] [llvm_offline_cache.cpp:load_meta_data@138] Lock /home/oliver/.cache/taichi/ticache/llvm/metadata.lock failed [Taichi] version 1.3.0, llvm 15.0.4, commit 0f25b95e, linux, python 3.10.8 [Taichi] Starting on arch=cuda [W 12/11/22 02:32:10.238 2903149] [offline_cache.h:run@176] Lock /home/oliver/.cache/taichi/ticache/llvm/metadata.lock failed

[Taichi] version 1.3.0, llvm 15.0.4, commit 0f25b95e, linux, python 3.10.8


Taichi Programming Language


Docs: https://docs.taichi-lang.org/ GitHub: https://github.com/taichi-dev/taichi/ Forum: https://forum.taichi.graphics/

                               TAICHI EXAMPLES                                    

──────────────────────────────────────────────────────────────────────────────────── 0: ad_gravity 24: laplace 48: physarum
1: comet 25: laplace_equation 49: print_offset
2: cornell_box 26: mandelbrot_zoom 50: rasterizer
3: diff_sph 27: marching_squares 51: regression
4: euler 28: mass_spring_3d_ggui 52: sdf_renderer
5: explicit_activation 29: mass_spring_game 53: simple_derivative
6: export_mesh 30: mass_spring_game_ggui 54: simple_texture
7: export_ply 31: mciso_advanced 55: simple_uv
8: export_videos 32: mgpcg 56: snow_phaseField
9: fem128 33: mgpcg_advanced 57: stable_fluid
10: fem128_ggui 34: minimal 58: stable_fluid_ggui
11: fem99 35: minimization 59: stable_fluid_graph
12: fractal 36: mpm128 60: taichi_bitmasked
13: fractal3d_ggui 37: mpm128_ggui 61: taichi_dynamic
14: fullscreen 38: mpm3d 62: taichi_logo
15: game_of_life 39: mpm3d_ggui 63: taichi_ngp
16: gui_image_io 40: mpm88 64: taichi_sparse
17: gui_widgets 41: mpm88_graph 65: texture_graph
18: implicit_fem 42: mpm99 66: tutorial
19: implicit_mass_spring 43: mpm_lagrangian_forces 67: two_stream_instability
20: initial_value_problem 44: nbody 68: vortex_rings
21: jacobian 45: odop_solar 69: waterwave
22: karman_vortex_street 46: patterns
23: keyboard 47: pbf2d
──────────────────────────────────────────────────────────────────────────────────── [W 12/11/22 02:32:12.044 2903194] [llvm_offline_cache.cpp:load_meta_data@138] Lock /home/oliver/.cache/taichi/ticache/llvm/metadata.lock failed 42 Running example minimal ... [Taichi] Starting on arch=x64 42.0

Running time: 0.58s [W 12/11/22 02:32:12.525 2903194] [offline_cache.h:run@176] Lock /home/oliver/.cache/taichi/ticache/llvm/metadata.lock failed [W 12/11/22 02:32:12.727 2903194] [llvm_offline_cache.cpp:dump@316] Lock /home/oliver/.cache/taichi/ticache/llvm/metadata.lock failed

Consider attaching this log when maintainers ask about system information.

Running time: 19.31s

oliver-batchelor commented 1 year ago

On a related point - I also can't pass my struct into the kernel.


@ti.kernel
def point_bounds(points:ti.template(), b:ti.template()):
  for i in points:
    b.lower = ti.math.min(b.lower, points[i])
    b.upper = ti.math.max(b.upper, points[i])
ti.init()

x = ti.Vector.field(3, dtype=ti.f32, shape=10)
points = np.arange(0, 30, dtype=np.float32).reshape(10, 3)
x.from_numpy(points)

b = AABox(vec3(0, 0, 0), vec3(0, 0, 0))
point_bounds(x, b)
print(b)

Hits another assert:

Traceback (most recent call last):
  File "/home/oliver/sync/python_geometry/test.py", line 25, in <module>
    point_bounds(x, b)
  File "/home/oliver/miniconda3/envs/torch13/lib/python3.10/site-packages/taichi/lang/kernel_impl.py", line 945, in wrapped
    return primal(*args, **kwargs)
  File "/home/oliver/miniconda3/envs/torch13/lib/python3.10/site-packages/taichi/lang/kernel_impl.py", line 871, in __call__
    key = self.ensure_compiled(*args)
  File "/home/oliver/miniconda3/envs/torch13/lib/python3.10/site-packages/taichi/lang/kernel_impl.py", line 836, in ensure_compiled
    instance_id, arg_features = self.mapper.lookup(args)
  File "/home/oliver/miniconda3/envs/torch13/lib/python3.10/site-packages/taichi/lang/kernel_impl.py", line 437, in lookup
    if key not in self.mapping:
TypeError: unhashable type: 'Struct'
bobcao3 commented 1 year ago

Sorry! This is something we are actively working on but not supported yet. @lin-hitonami has it supported for real_functions and we are trying to get struct type into the actual type system.

oliver-batchelor commented 1 year ago

@bobcao3 Ah - thank you!