SVE Instruction to get index of first true element

Hi Maxbit,

You are correct that there is no way to directly extract the index of the first predicated element in SVE, however there are a couple of different approaches you can take to achieve this.

In particular you can use the `svindex` intrinsic to generate a vector of 0,1,... and then use `svlastb` on that instead, which will give you the index you are after.

It is also worth pointing out that your question says the first matching element, but `svlastb` will return the last matching element instead. If you can guarantee that there is only a single matching element then this is fine, but if not then you may need to first adjust the predicate such that the last element is the one you are after. You can do this with either `svbrka_b_z` or `svpfirst_b` to ensure that the first matching element of your input predicate is the last one of the resulting predicate that you pass to `svlastb`.

Putting that all together and you might end up with something like this:

#include <arm_sve.h>
#include <stdio.h>

uint64_t find_first_index(svuint64_t x, svuint64_t y) {
  svbool_t ptrue = svptrue_b8();
  svbool_t p = svcmpeq_u64(ptrue, x, y);
  return svlastb_u64(svbrka_b_z(p, p), svindex_u64(0, 1));
}

int main() {
  uint64_t data1[4] = { 0, 1, 2, 3 }; // assume vl=256, adjust as needed
  uint64_t data2[4] = { 3, 2, 2, 3 };
  svbool_t ptrue = svptrue_b8();
  svuint64_t x = svld1(ptrue, data1);
  svuint64_t y = svld1(ptrue, data2);
  printf("%lu\n", find_first_index(x, y));
}

You will need to be careful that you do not try to use this in a situation where none of the elements match, but you can avoid this by simply checking `svptest_any` and continuing your search if none are found.

Hope that helps,
George

Hi Maxbit,

You are correct that there is no way to directly extract the index of the first predicated element in SVE, however there are a couple of different approaches you can take to achieve this.

In particular you can use the `svindex` intrinsic to generate a vector of 0,1,... and then use `svlastb` on that instead, which will give you the index you are after.

It is also worth pointing out that your question says the first matching element, but `svlastb` will return the last matching element instead. If you can guarantee that there is only a single matching element then this is fine, but if not then you may need to first adjust the predicate such that the last element is the one you are after. You can do this with either `svbrka_b_z` or `svpfirst_b` to ensure that the first matching element of your input predicate is the last one of the resulting predicate that you pass to `svlastb`.

Putting that all together and you might end up with something like this:

#include <arm_sve.h>
#include <stdio.h>

uint64_t find_first_index(svuint64_t x, svuint64_t y) {
  svbool_t ptrue = svptrue_b8();
  svbool_t p = svcmpeq_u64(ptrue, x, y);
  return svlastb_u64(svbrka_b_z(p, p), svindex_u64(0, 1));
}

int main() {
  uint64_t data1[4] = { 0, 1, 2, 3 }; // assume vl=256, adjust as needed
  uint64_t data2[4] = { 3, 2, 2, 3 };
  svbool_t ptrue = svptrue_b8();
  svuint64_t x = svld1(ptrue, data1);
  svuint64_t y = svld1(ptrue, data2);
  printf("%lu\n", find_first_index(x, y));
}

You will need to be careful that you do not try to use this in a situation where none of the elements match, but you can avoid this by simply checking `svptest_any` and continuing your search if none are found.

Hope that helps,
George

(a) For comparing different choices of instruction sequences you can refer to the software optimisation guides available on developer.arm.com, for example for Neoverse V1: https://developer.arm.com/documentation/pjdoc466751330-9685/latest/ . On Neoverse V1 at least it appears that BRKA and PFIRST instructions have identical latency and throughput, so which one you use probably doesn't matter, but I would suggest checking the guide for the exact core you are interested in.

(b) For the Neon equivalent as you point out there is no nice sequence here. The fastest sequence will probably depend both on how frequently the INVALID_IDX path is taken compared to an actual index being found, as well as the particular core you are interested in, so I'd recomment benchmarking a couple of different alternatives yourself to see which one works better for your use case.

For the code itself: comparisons in Neon like `vceqq` will set matching elements to all-ones, so you can take advantage of that in a couple of ways:

• You can use the UMAXV instruction and then compare this against zero to see if any elements match.
• You can use the comparison result and AND that against a hard-coded vector of {1,2,...}, then use UMINV, take one away from that to find the minimum matching index (you need to avoid zero here since the non-matching cases will already be zero). Unlike SVE you don't need an INDEX instruction here, you can just hard-code it since the vector length is known at compile time.
• You can alternatively use BSL instead of AND to select the passing elements into a {0,1,...} vector and the others to -1, to avoid needing to subtract one at the end.

Since Neon vectors are only a pair of uint64_t elements wide, you might actually find it quicker to extract out the pair of elements and just call e.g. __builtin_ctz twice. In general the scalar instructions have lower overheads for such short sequences.

I appreciate that's a bit of a vague answer. Let me know if you run into any other problems and I can try and throw a more complete example together.