Issue with tile assignment involving composite types

[shi-eric]

Potentially related:

import numpy as np
import warp as wp

vec256f = wp.types.vector(length=256, dtype=wp.float32)

@wp.kernel
def test(a: wp.array2d(dtype=wp.float32)):
    i = wp.tid()
    R = wp.tile_zeros(shape=(1,), dtype=vec256f)
    R[0][i] = wp.float32(i)
    a[i // 16, i % 16] = R[0][i]  # incorrect
    # a[i // 16, i % 16] = wp.float32(i)  # correct

TILE_SIZE = 256
device="cuda"
a = wp.zeros((16, 16), dtype=wp.float32, device=device)
wp.launch(kernel=test, dim=TILE_SIZE, outputs=[a], block_dim=TILE_SIZE, device=device)
wp.synchronize()
print(a.numpy())

kernel: test dim: 256 inputs: [] outputs: [array(shape=(16, 16), dtype=float32)] device: cuda:0
[[0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]]

I'll try to test this on v1.9.0 later.

Originally posted by @j3soon in #926

[j3soon]

Currently, the code will have runtime error on latest commit on main: 102d8a7.

I can reproduce this issue (no runtime error, but incorrect result) on commit f05553a (i.e., parent commit of 195f774). However, after commit 195f774, the runtime error occurs. I suspect this is potentially caused by interpreting R[0][i] as R[0, i] for all tile arrays in commit 195f774 in order to fix #793 and #926?

We may need to consider that for a tile array R with data type vector/matrix, the elements can only be referred to as R[arr_idx][vec_idx] and not R[arr_idx][vec_idx].

Adding @c0d1f1ed for viz

Full runtime error output:

Warp 1.10.0.dev0 initialized:
   CUDA Toolkit 13.0, Driver 13.0
   Devices:
     "cpu"      : "x86_64"
     "cuda:0"   : "NVIDIA L40" (47 GiB, sm_89, mempool enabled)
   Kernel cache:
     /root/.cache/warp/1.10.0.dev0
Module __main__ 0d86fe9 load on device 'cuda:0' took 0.82 ms  (error)
Traceback (most recent call last):
  File "/root/test.py", line 17, in <module>
    wp.launch(kernel=test, dim=TILE_SIZE, outputs=[a], block_dim=TILE_SIZE, device=device)
    ~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/venv/lib/python3.13/site-packages/warp/context.py", line 6171, in launch
    module_exec = kernel.module.load(device, block_dim)
  File "/opt/venv/lib/python3.13/site-packages/warp/context.py", line 2601, in load
    raise (e)
  File "/opt/venv/lib/python3.13/site-packages/warp/context.py", line 2598, in load
    self.compile(device, module_dir, output_name, output_arch)
    ~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/venv/lib/python3.13/site-packages/warp/context.py", line 2371, in compile
    builder = ModuleBuilder(
        self,
        builder_options,
        hasher=self.hashers.get(self.options["block_dim"], None),
    )
  File "/opt/venv/lib/python3.13/site-packages/warp/context.py", line 1784, in __init__
    self.build_kernel(kernel)
    ~~~~~~~~~~~~~~~~~^^^^^^^^
  File "/opt/venv/lib/python3.13/site-packages/warp/context.py", line 1816, in build_kernel
    kernel.adj.build(self)
    ~~~~~~~~~~~~~~~~^^^^^^
  File "/opt/venv/lib/python3.13/site-packages/warp/codegen.py", line 1100, in build
    raise type(original_exc)(*new_args).with_traceback(original_exc.__traceback__) from None
  File "/opt/venv/lib/python3.13/site-packages/warp/codegen.py", line 1088, in build
    adj.eval(adj.tree.body[0])
    ~~~~~~~~^^^^^^^^^^^^^^^^^^
  File "/opt/venv/lib/python3.13/site-packages/warp/codegen.py", line 3063, in eval
    return emit_node(adj, node)
  File "/opt/venv/lib/python3.13/site-packages/warp/codegen.py", line 1773, in emit_FunctionDef
    adj.eval(f)
    ~~~~~~~~^^^
  File "/opt/venv/lib/python3.13/site-packages/warp/codegen.py", line 3063, in eval
    return emit_node(adj, node)
  File "/opt/venv/lib/python3.13/site-packages/warp/codegen.py", line 2654, in emit_Assign
    rhs = adj.eval(node.value)
  File "/opt/venv/lib/python3.13/site-packages/warp/codegen.py", line 3063, in eval
    return emit_node(adj, node)
  File "/opt/venv/lib/python3.13/site-packages/warp/codegen.py", line 2626, in emit_Subscript
    return adj.emit_indexing(target, indices)
           ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/opt/venv/lib/python3.13/site-packages/warp/codegen.py", line 2546, in emit_indexing
    raise RuntimeError(
        f"Incorrect number of indices specified for a tile view/extract, got {len(indices)} indices for a {len(target_type.shape)} dimensional tile."
    )
RuntimeError: Error while parsing function "test" at /root/test.py:11:
    a[i // 16, i % 16] = R[0][i]  # incorrect
;Incorrect number of indices specified for a tile view/extract, got 2 indices for a 1 dimensional tile.

[j3soon]

May be related to this snippet removed in commit 195f774:

warp/warp/codegen.py

Lines 2464 to 2515 in f05553a

	# returns the object being indexed, and the list of indices
	def eval_subscript(adj, node):
	# We want to coalesce multi-dimensional array indexing into a single operation. This needs to deal with expressions like `a[i][j][x][y]` where `a` is a 2D array of matrices,
	# and essentially rewrite it into `a[i, j][x][y]`. Since the AST observes the indexing right-to-left, and we don't want to evaluate the index expressions prematurely,
	# this requires a first loop to check if this `node` only performs indexing on the array, and a second loop to evaluate and collect index variables.
	root = node
	count = 0
	array = None
	while isinstance(root, ast.Subscript):
	if isinstance(root.slice, ast.Tuple):
	# handles the x[i, j] case (Python 3.8.x upward)
	count += len(root.slice.elts)
	elif isinstance(root.slice, ast.Index) and isinstance(root.slice.value, ast.Tuple):
	# handles the x[i, j] case (Python 3.7.x)
	count += len(root.slice.value.elts)
	else:
	# simple expression, e.g.: x[i]
	count += 1
	if isinstance(root.value, ast.Name):
	symbol = adj.emit_Name(root.value)
	symbol_type = strip_reference(symbol.type)
	if is_array(symbol_type):
	array = symbol
	break
	root = root.value
	# If not all indices index into the array, just evaluate the right-most indexing operation.
	if not array or (count > array.type.ndim):
	count = 1
	nodes = ()
	root = node
	while len(nodes) < count:
	if isinstance(root.slice, ast.Tuple):
	ij = tuple(arg for arg in root.slice.elts)
	elif isinstance(root.slice, ast.Index) and isinstance(root.slice.value, ast.Tuple):
	# The node `ast.Index` is deprecated in Python 3.9.
	ij = tuple(arg for arg in root.slice.value.elts)
	elif isinstance(root.slice, ast.ExtSlice):
	# The node `ast.ExtSlice` is deprecated in Python 3.9.
	ij = tuple(arg for arg in root.slice.dims)
	else:
	ij = (root.slice,)
	nodes = ij + nodes # prepend
	root = root.value
	target = adj.eval(root)
	target_type = strip_reference(target.type)