Skip to content

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

@hakanai

Description

@hakanai

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".

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:

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    apiCommon apienhancementNew feature or request

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions