Some efficient method for permuting elements of an NDArray to get a new NDArray
- 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
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