Kotlin / Kotlin/multik

Some efficient method for permuting elements of an NDArray to get a new NDArray

Open
#215 1 comment 0 reactions 0 assignees View on GitHub
api feature ndarray perf
Dominant language
Kotlin
Stars
734
Forks
51
Avg merge
9h 52m
Merged PRs (30d)
6

Description

One of the things we have in Vector API and which is also very common in shader code, is the ability to permute the indices of a vector to get a new vector. In computer graphics, we tend to call this "swizzling".

```glsl
float4 a = float4(1, 2, 3, 4);
float4 b = a.wwxy; // b now contains {4,4,1,2}
```

One use case I found this to help was vector cross products. Of course, in shader code, you have the cross product as a primitive, but not so lucky in Kotlin, so I hacked something up for Vector API, working something like this:

```kotlin
import jdk.incubator.vector.DoubleVector
import jdk.incubator.vector.VectorShuffle

private val yzxwShuffle = VectorShuffle.fromValues(DoubleVector.SPECIES_256, 1, 2, 0, 3)

private val zxywShuffle = VectorShuffle.fromValues(DoubleVector.SPECIES_256, 2, 0, 1, 3)

val DoubleVector.yzxw: DoubleVector
get() = rearrange(yzxwShuffle)

val DoubleVector.zxyw: DoubleVector
get() = rearrange(zxywShuffle)

// implementing vector cross product
val resultCells = (cells.yzxw * their.cells.zxyw) - (cells.zxyw * their.cells.yzxw)
```

I'm not sure if there is a better performance option for this one, but I was pretty happy with how I got to use the same pair of swizzles for both sides of the operation, just swapping which cells they were being applied to. I assume that these vector shuffle operations are fairly cheap to perform even on CPUs, because I was able to speed up my cross product quite a bit using this.

So I would like access to something similar in `NDArray`, or at least in `D1Array`. I don't have a good vision of what something like this would look like for higher dimension arrays. I thought I might be able to use `slice()` to get what I want, but it looks like each dimension you pass to that method can only return either one cell or a range of cells. I would like to pass a list of indices which may not be in order.

Contributor guide

Open the contributing guide

Research direction

Start by examining the existing NDArray and D1Array slice() APIs and how they handle per-dimension indices. Define how unordered index lists should produce a new array, including whether the feature applies only to D1Array or higher-dimensional arrays, then verify the chosen behavior with focused coverage.

Written by the indexing model from the issue text.

Assessment

Tech stack
kotlin
Domain
data
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.