NVIDIA / NVIDIA/cccl

Allow unpacking of tuple in cuda.compute operators when a `ZipIterator` is used as input

Open
#6,610 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
2.5k
Forks
486
Avg merge
2d 6h
Merged PRs (30d)
295

Description

It is currently not possible to unpack values from a tuple in an operator we use in cuda.compute due to a numba compilation error:

```python
def op(x, y):
a, b = x
# ...
```

```
E During: Pass nopython_type_inference
E During: lowering "$20call.5 = call $4load_global.0(arg_0, arg_1, ret, func=$4load_global.0, args=[Var(arg_0, :27), Var(arg_1, :27), Var(ret, :27)], kws=(), vararg=None, varkwarg=None, target=None)" at (28)
E During: Pass cuda_native_lowering
```

In #6518, we initially implemented this by adding this overload to `_setup_numba_struct`:

```python
# Add tuple() conversion support to enable unpacking with explicit conversion
# Usage: a, b = tuple(struct)
@overload(tuple)
def struct_as_tuple(struct_val):
if not isinstance(struct_val, StructType):
return

# Generate code to extract all fields and return as tuple
field_accesses = [f"struct_val.{fname}" for fname in field_names_list]
impl_code = f"""
def struct_as_tuple_impl(struct_val):
return ({", ".join(field_accesses)},)
"""
local_vars = {}
exec(impl_code, {}, local_vars)
return local_vars["struct_as_tuple_impl"]
```

However this required us to explicitly construct a tuple:

```python
def op(x, y):
a, b = tuple(x)
# ...
```

So we removed it. We would still like to be able to unpack a tuple without explicitly using the constructor.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.