# Idioms for using an index tensor in kernel computations

**URL:** <https://forum.modular.com/t/idioms-for-using-an-index-tensor-in-kernel-computations/1593>\
**Category:** GPU Programming\
**Created:** [June 6, 2025, 3:33pm UTC](https://forum.modular.com/t/idioms-for-using-an-index-tensor-in-kernel-computations/1593 "2025-06-06T15:33:27Z")\
**Posts on this page:** 8\
**Page:** 1

<div class="post-metadata">

**Author:** ![femtomc](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/femtomc/32/542_2.png) [@femtomc](https://forum.modular.com/u/femtomc)\
**Post date:** [June 6, 2025, 3:33pm UTC](https://forum.modular.com/t/idioms-for-using-an-index-tensor-in-kernel-computations/1593/1 "2025-06-06T15:33:27Z")

</div>

Hi all!

I’m trying to express a particular idiom within a GPU kernel, where I grab an index out of a tensor, and then use it to index into another tensor.

The full code is below for reference, but below the code I show the relevant segment.

```mojo
# GPU kernel for computing cross-covariance matrix elements
fn compute_cross_covariance_kernel[
    n: Int
](
    W_tensor: LayoutTensor[mut=True, float32, Layout.row_major(3, 3)],
    correspondences_tensor: LayoutTensor[
        mut=True, DType.int32, Layout.row_major(n)
    ],
    distances_tensor: LayoutTensor[mut=True, float32, Layout.row_major(n)],
    source_centered_tensor: LayoutTensor[
        mut=False, float32, Layout.row_major(n, 3)
    ],
    target_centered_tensor: LayoutTensor[
        mut=False, float32, Layout.row_major(n, 3)
    ],
    max_dist: Float32,
):
    # Each block computes one element of the 3x3 matrix
    row = block_idx.x
    col = block_idx.y

    if row >= 3 or col >= 3:
        return

    # Shared stack memory for partial sums
    var shared = stack_allocation[
        THREADS_PER_BLOCK,
        Scalar[float32],
        address_space = AddressSpace.SHARED,
    ]()

    tid = thread_idx.x
    var sum = 0.0

    # Each thread processes multiple correspondences
    var idx = tid
    while idx < n:
        # Check if correspondence is valid (within distance threshold)
        if distances_tensor[idx] < max_dist:
            var target_idx = correspondences_tensor[idx]
            s_val = source_centered_tensor[idx, row]
            t_val = target_centered_tensor[target_idx, col]
            sum += s_val * t_val

        idx += THREADS_PER_BLOCK

    # Store in shared memory and reduce
    shared[tid] = sum
    barrier()

    # Parallel reduction
    var stride = THREADS_PER_BLOCK // 2
    while stride > 0:
        if tid < stride:
            shared[tid] = shared[tid] + shared[tid + stride]
        barrier()
        stride //= 2

    # Thread 0 writes final result
    if tid == 0:
        W_tensor[row, col] = shared[0]

```

The relevant part is here:

```auto
var target_idx: Int = correspondences_tensor[idx]
s_val = source_centered_tensor[idx, row]
t_val = target_centered_tensor[target_idx, col]
sum += s_val * t_val

```

trying this _as-is_ runs into a resolution error (as far as I can tell).

 ![image](https://us1.discourse-cdn.com/flex001/uploads/modular/original/1X/12d4c66adf0e6728dcbec2743abed85306a7ff0a.png)

Is there a known idiom for this? I did a cursory search through the GPU puzzles but not experienced enough to find a good match.

---

<div class="post-metadata">

**Author:** ![lukas](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/lukas/32/47_2.png) [@lukas](https://forum.modular.com/u/lukas)\
**Post date:** [June 6, 2025, 4:32pm UTC](https://forum.modular.com/t/idioms-for-using-an-index-tensor-in-kernel-computations/1593/2 "2025-06-06T16:32:53Z")

</div>

The issue here is that `LayoutTensor. __getitem__ ` returns a `SIMD[dtype, Self.element_size]` (basically a Tuple of size 1 in your case). You can see that from the error message (though we should clean that up a little). If you wrote `correspondences_tensor[idx][0]` or `Int(correspondences_tensor[idx])` it would work as an index. We are actively exploring how to make that API better.

---

<div class="post-metadata">

**Author:** ![femtomc](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/femtomc/32/542_2.png) [@femtomc](https://forum.modular.com/u/femtomc)\
**Post date:** [June 6, 2025, 4:33pm UTC](https://forum.modular.com/t/idioms-for-using-an-index-tensor-in-kernel-computations/1593/3 "2025-06-06T16:33:35Z")

</div>

Doing some `rebind` hacking allows me to compile, but running into

```auto
double free or corruption (fasttop)

```

when attempting to run. Debugging that now (assuming that my usage of `rebind` is okay):

```auto
target_idx = rebind[Int](correspondences_tensor[idx])
s_val = source_centered_tensor[idx, row]
t_val = target_centered_tensor[target_idx, col]
sum += rebind[Float32](s_val * t_val)

```

Edit: @lukas fix above worked, and removed my memory problem.

---

<div class="post-metadata">

**Author:** ![femtomc](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/femtomc/32/542_2.png) [@femtomc](https://forum.modular.com/u/femtomc)\
**Post date:** [June 6, 2025, 4:36pm UTC](https://forum.modular.com/t/idioms-for-using-an-index-tensor-in-kernel-computations/1593/4 "2025-06-06T16:36:01Z")

</div>

Surprised to find that `Int(...)` works, but `Float32(...)` doesn’t when used in the same capacity:

 ![image](https://us1.discourse-cdn.com/flex001/uploads/modular/original/1X/da1f1c0f6fdaba95a1cc70d0acf02780cd3a7a12.png)

requiring a `rebind` (to work)

---

<div class="post-metadata">

**Author:** ![lukas](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/lukas/32/47_2.png) [@lukas](https://forum.modular.com/u/lukas)\
**Post date:** [June 6, 2025, 4:37pm UTC](https://forum.modular.com/t/idioms-for-using-an-index-tensor-in-kernel-computations/1593/5 "2025-06-06T16:37:16Z")

</div>

If you think about it, `Float32(source_centered_tensor[Idx, row])` is writing:

```mojo
SIMD[DType.float32, 1](SIMD[DType.float32, N](...))

```

which doesn’t really make sense

---

<div class="post-metadata">

**Author:** ![lukas](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/lukas/32/47_2.png) [@lukas](https://forum.modular.com/u/lukas)\
**Post date:** [June 6, 2025, 4:38pm UTC](https://forum.modular.com/t/idioms-for-using-an-index-tensor-in-kernel-computations/1593/6 "2025-06-06T16:38:16Z")

</div>

`Int` isn’t an alias for `SIMD`, it’s a unique type closer to `isize` in Rust (or `ssize_t` in C++)

---

<div class="post-metadata">

**Author:** ![femtomc](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/femtomc/32/542_2.png) [@femtomc](https://forum.modular.com/u/femtomc)\
**Post date:** [June 6, 2025, 4:39pm UTC](https://forum.modular.com/t/idioms-for-using-an-index-tensor-in-kernel-computations/1593/7 "2025-06-06T16:39:27Z")

</div>

Excellent, thanks – was just about to ask what the difference between `Int` and `Float32` is.

---

<div class="post-metadata">

**Author:** ![lukas](https://sea1.discourse-cdn.com/flex001/user_avatar/forum.modular.com/lukas/32/47_2.png) [@lukas](https://forum.modular.com/u/lukas)\
**Post date:** [June 6, 2025, 4:40pm UTC](https://forum.modular.com/t/idioms-for-using-an-index-tensor-in-kernel-computations/1593/8 "2025-06-06T16:40:19Z")

</div>

That being said, rather than `rebind`, you can write `source_centered_tensor[idx, row][0]`