You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by ka...@apache.org on 2026/02/26 16:38:30 UTC
(lucene) branch main updated: Optimize vectorized int4{DotProduct,SquareDistance} computations by avoiding conversions (#15742)
This is an automated email from the ASF dual-hosted git repository.
kaivalnp pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/lucene.git
The following commit(s) were added to refs/heads/main by this push:
new 98e1c075f70 Optimize vectorized int4{DotProduct,SquareDistance} computations by avoiding conversions (#15742)
98e1c075f70 is described below
commit 98e1c075f709c13d501f91091e1b05f032279e48
Author: Kaival Parikh <ka...@gmail.com>
AuthorDate: Thu Feb 26 11:38:24 2026 -0500
Optimize vectorized int4{DotProduct,SquareDistance} computations by avoiding conversions (#15742)
Replace intermediate vector conversions with reinterpret casting + bit manipulation.
---
lucene/CHANGES.txt | 2 +
.../vectorization/PanamaVectorUtilSupport.java | 169 +++++++++++++--------
2 files changed, 105 insertions(+), 66 deletions(-)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index 272dd0210c6..c2cd2bac185 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -233,6 +233,8 @@ Optimizations
* GITHUB#15729: Lucene90DocValuesProducer is not prefetching any data for DocValueSkippers anymore (Alexander Reelsen)
+* GITHUB#15742: Optimize int4 dotProduct and squareDistance computations by replacing vector conversions with reinterpret casting + bit manipulation. (Trevor McCulloch, Kaival Parikh)
+
Bug Fixes
---------------------
* GITHUB#15668: Fix UnsupportedOperationException in WeightedSpanTermExtractor by
diff --git a/lucene/core/src/java25/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java b/lucene/core/src/java25/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java
index 6f0407f1d66..d41118a49b0 100644
--- a/lucene/core/src/java25/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java
+++ b/lucene/core/src/java25/org/apache/lucene/internal/vectorization/PanamaVectorUtilSupport.java
@@ -490,17 +490,17 @@ final class PanamaVectorUtilSupport implements VectorUtilSupport {
static {
if (VECTOR_BITSIZE >= 512) {
- BYTE_SPECIES = ByteVector.SPECIES_256;
+ BYTE_SPECIES = ByteVector.SPECIES_512;
SHORT_SPECIES = ShortVector.SPECIES_512;
- CHUNK = 4096;
+ CHUNK = 8192;
} else if (VECTOR_BITSIZE == 256) {
- BYTE_SPECIES = ByteVector.SPECIES_128;
+ BYTE_SPECIES = ByteVector.SPECIES_256;
SHORT_SPECIES = ShortVector.SPECIES_256;
- CHUNK = 2048;
+ CHUNK = 4096;
} else {
- BYTE_SPECIES = ByteVector.SPECIES_64;
+ BYTE_SPECIES = ByteVector.SPECIES_128;
SHORT_SPECIES = ShortVector.SPECIES_128;
- CHUNK = 1024;
+ CHUNK = 2048;
}
}
}
@@ -536,22 +536,30 @@ final class PanamaVectorUtilSupport implements VectorUtilSupport {
int sum = 0;
// iterate in chunks to ensure we don't overflow the short accumulator
for (int i = 0; i < limit; i += Int4Constants.CHUNK) {
- ShortVector acc = ShortVector.zero(Int4Constants.SHORT_SPECIES);
+ ShortVector acc0 = ShortVector.zero(Int4Constants.SHORT_SPECIES);
+ ShortVector acc1 = ShortVector.zero(Int4Constants.SHORT_SPECIES);
int innerLimit = Math.min(limit - i, Int4Constants.CHUNK);
for (int j = 0; j < innerLimit; j += Int4Constants.BYTE_SPECIES.length()) {
// unpacked
ByteVector vb8 = b.load(Int4Constants.BYTE_SPECIES, i + j);
- Vector<Short> vb16 = vb8.convertShape(B2S, Int4Constants.SHORT_SPECIES, 0);
// unpacked
ByteVector va8 = a.load(Int4Constants.BYTE_SPECIES, i + j);
- Vector<Short> va16 = va8.convertShape(B2S, Int4Constants.SHORT_SPECIES, 0);
- acc = acc.add(vb16.mul(va16));
+ ShortVector prod16 = va8.mul(vb8).reinterpretAsShorts();
+ acc0 = acc0.add(prod16.lanewise(LSHR, 8));
+ acc1 = acc1.add(prod16.and((short) 0xFF));
}
- Vector<Integer> intAcc0 = acc.convert(S2I, 0);
- Vector<Integer> intAcc1 = acc.convert(S2I, 1);
- sum += intAcc0.add(intAcc1).reinterpretAsInts().reduceLanes(ADD);
+
+ IntVector intAcc0 = acc0.reinterpretAsInts();
+ IntVector intAcc1 = acc1.reinterpretAsInts();
+ sum +=
+ intAcc0
+ .and(0xFFFF)
+ .add(intAcc0.lanewise(LSHR, 16))
+ .add(intAcc1.and(0xFFFF))
+ .add(intAcc1.lanewise(LSHR, 16))
+ .reduceLanes(ADD);
}
return sum;
}
@@ -597,24 +605,28 @@ final class PanamaVectorUtilSupport implements VectorUtilSupport {
// packed
ByteVector vb8 = packed.load(Int4Constants.BYTE_SPECIES, i + j);
+ // lower
+ ByteVector vc8 = unpacked.load(Int4Constants.BYTE_SPECIES, i + j);
+
// upper
ByteVector va8 = unpacked.load(Int4Constants.BYTE_SPECIES, i + j + packed.length());
- ByteVector prod8 = vb8.and((byte) 0x0F).mul(va8);
- Vector<Short> prod16 = prod8.convertShape(ZERO_EXTEND_B2S, Int4Constants.SHORT_SPECIES, 0);
- acc0 = acc0.add(prod16);
- // lower
- ByteVector vc8 = unpacked.load(Int4Constants.BYTE_SPECIES, i + j);
- ByteVector prod8a = vb8.lanewise(LSHR, 4).mul(vc8);
- Vector<Short> prod16a =
- prod8a.convertShape(ZERO_EXTEND_B2S, Int4Constants.SHORT_SPECIES, 0);
- acc1 = acc1.add(prod16a);
+ ShortVector prod16 = vb8.and((byte) 0x0F).mul(va8).reinterpretAsShorts();
+ acc0 = acc0.add(prod16.lanewise(LSHR, 8).add(prod16.and((short) 0xFF)));
+
+ ShortVector prod16a = vb8.lanewise(LSHR, 4).mul(vc8).reinterpretAsShorts();
+ acc1 = acc1.add(prod16a.lanewise(LSHR, 8).add(prod16a.and((short) 0xFF)));
}
- Vector<Integer> intAcc0 = acc0.convert(S2I, 0);
- Vector<Integer> intAcc1 = acc0.convert(S2I, 1);
- Vector<Integer> intAcc2 = acc1.convert(S2I, 0);
- Vector<Integer> intAcc3 = acc1.convert(S2I, 1);
- sum += intAcc0.add(intAcc1).add(intAcc2).add(intAcc3).reinterpretAsInts().reduceLanes(ADD);
+
+ IntVector intAcc0 = acc0.reinterpretAsInts();
+ IntVector intAcc1 = acc1.reinterpretAsInts();
+ sum +=
+ intAcc0
+ .and(0xFFFF)
+ .add(intAcc0.lanewise(LSHR, 16))
+ .add(intAcc1.and(0xFFFF))
+ .add(intAcc1.lanewise(LSHR, 16))
+ .reduceLanes(ADD);
}
return sum;
}
@@ -660,21 +672,24 @@ final class PanamaVectorUtilSupport implements VectorUtilSupport {
var va8 = a.load(Int4Constants.BYTE_SPECIES, i + j);
// upper
- ByteVector prod8 = vb8.and((byte) 0x0F).mul(va8.and((byte) 0x0F));
- Vector<Short> prod16 = prod8.convertShape(ZERO_EXTEND_B2S, Int4Constants.SHORT_SPECIES, 0);
- acc0 = acc0.add(prod16);
+ ShortVector prod16 = vb8.and((byte) 0x0F).mul(va8.and((byte) 0x0F)).reinterpretAsShorts();
+ acc0 = acc0.add(prod16.lanewise(LSHR, 8).add(prod16.and((short) 0xFF)));
// lower
- ByteVector prod8a = vb8.lanewise(LSHR, 4).mul(va8.lanewise(LSHR, 4));
- Vector<Short> prod16a =
- prod8a.convertShape(ZERO_EXTEND_B2S, Int4Constants.SHORT_SPECIES, 0);
- acc1 = acc1.add(prod16a);
+ ShortVector prod16a =
+ vb8.lanewise(LSHR, 4).mul(va8.lanewise(LSHR, 4)).reinterpretAsShorts();
+ acc1 = acc1.add(prod16a.lanewise(LSHR, 8).add(prod16a.and((short) 0xFF)));
}
- Vector<Integer> intAcc0 = acc0.convert(S2I, 0);
- Vector<Integer> intAcc1 = acc0.convert(S2I, 1);
- Vector<Integer> intAcc2 = acc1.convert(S2I, 0);
- Vector<Integer> intAcc3 = acc1.convert(S2I, 1);
- sum += intAcc0.add(intAcc1).add(intAcc2).add(intAcc3).reinterpretAsInts().reduceLanes(ADD);
+
+ IntVector intAcc0 = acc0.reinterpretAsInts();
+ IntVector intAcc1 = acc1.reinterpretAsInts();
+ sum +=
+ intAcc0
+ .and(0xFFFF)
+ .add(intAcc0.lanewise(LSHR, 16))
+ .add(intAcc1.and(0xFFFF))
+ .add(intAcc1.lanewise(LSHR, 16))
+ .reduceLanes(ADD);
}
return sum;
}
@@ -946,7 +961,8 @@ final class PanamaVectorUtilSupport implements VectorUtilSupport {
int sum = 0;
// iterate in chunks to ensure we don't overflow the short accumulator
for (int i = 0; i < limit; i += Int4Constants.CHUNK) {
- ShortVector acc = ShortVector.zero(Int4Constants.SHORT_SPECIES);
+ ShortVector acc0 = ShortVector.zero(Int4Constants.SHORT_SPECIES);
+ ShortVector acc1 = ShortVector.zero(Int4Constants.SHORT_SPECIES);
int innerLimit = Math.min(limit - i, Int4Constants.CHUNK);
for (int j = 0; j < innerLimit; j += Int4Constants.BYTE_SPECIES.length()) {
// unpacked
@@ -954,13 +970,22 @@ final class PanamaVectorUtilSupport implements VectorUtilSupport {
// unpacked
var va8 = a.load(Int4Constants.BYTE_SPECIES, i + j);
- ByteVector diff8 = vb8.sub(va8);
- Vector<Short> diff16 = diff8.convertShape(B2S, Int4Constants.SHORT_SPECIES, 0);
- acc = acc.add(diff16.mul(diff16));
+ ShortVector diff8 = vb8.sub(va8).abs().reinterpretAsShorts();
+ ShortVector diff16 = diff8.and((short) 0xFF);
+ acc0 = acc0.add(diff16.mul(diff16));
+ ShortVector diff16a = diff8.lanewise(LSHR, 8);
+ acc1 = acc1.add(diff16a.mul(diff16a));
}
- Vector<Integer> intAcc0 = acc.convert(S2I, 0);
- Vector<Integer> intAcc1 = acc.convert(S2I, 1);
- sum += intAcc0.add(intAcc1).reinterpretAsInts().reduceLanes(ADD);
+
+ IntVector intAcc0 = acc0.reinterpretAsInts();
+ IntVector intAcc1 = acc1.reinterpretAsInts();
+ sum +=
+ intAcc0
+ .and(0xFFFF)
+ .add(intAcc0.lanewise(LSHR, 16))
+ .add(intAcc1.and(0xFFFF))
+ .add(intAcc1.lanewise(LSHR, 16))
+ .reduceLanes(ADD);
}
return sum;
}
@@ -1008,23 +1033,30 @@ final class PanamaVectorUtilSupport implements VectorUtilSupport {
// packed
ByteVector vb8 = packed.load(Int4Constants.BYTE_SPECIES, i + j);
+ // lower
+ ByteVector vc8 = unpacked.load(Int4Constants.BYTE_SPECIES, i + j);
+
// upper
ByteVector va8 = unpacked.load(Int4Constants.BYTE_SPECIES, i + j + packed.length());
+
ByteVector diff8 = vb8.and((byte) 0x0F).sub(va8);
- Vector<Short> diff16 = diff8.convertShape(B2S, Int4Constants.SHORT_SPECIES, 0);
- acc0 = acc0.add(diff16.mul(diff16));
+ ShortVector prod16 = diff8.mul(diff8).reinterpretAsShorts();
+ acc0 = acc0.add(prod16.and((short) 0xFF).add(prod16.lanewise(LSHR, 8)));
- // lower
- ByteVector vc8 = unpacked.load(Int4Constants.BYTE_SPECIES, i + j);
ByteVector diff8a = vb8.lanewise(LSHR, 4).sub(vc8);
- Vector<Short> diff16a = diff8a.convertShape(B2S, Int4Constants.SHORT_SPECIES, 0);
- acc1 = acc1.add(diff16a.mul(diff16a));
+ ShortVector prod16a = diff8a.mul(diff8a).reinterpretAsShorts();
+ acc1 = acc1.add(prod16a.and((short) 0xFF).add(prod16a.lanewise(LSHR, 8)));
}
- Vector<Integer> intAcc0 = acc0.convert(S2I, 0);
- Vector<Integer> intAcc1 = acc0.convert(S2I, 1);
- Vector<Integer> intAcc2 = acc1.convert(S2I, 0);
- Vector<Integer> intAcc3 = acc1.convert(S2I, 1);
- sum += intAcc0.add(intAcc1).add(intAcc2).add(intAcc3).reinterpretAsInts().reduceLanes(ADD);
+
+ IntVector intAcc0 = acc0.reinterpretAsInts();
+ IntVector intAcc1 = acc1.reinterpretAsInts();
+ sum +=
+ intAcc0
+ .and(0xFFFF)
+ .add(intAcc0.lanewise(LSHR, 16))
+ .add(intAcc1.and(0xFFFF))
+ .add(intAcc1.lanewise(LSHR, 16))
+ .reduceLanes(ADD);
}
return sum;
}
@@ -1074,19 +1106,24 @@ final class PanamaVectorUtilSupport implements VectorUtilSupport {
// upper
ByteVector diff8 = vb8.and((byte) 0x0F).sub(va8.and((byte) 0x0F));
- Vector<Short> diff16 = diff8.convertShape(B2S, Int4Constants.SHORT_SPECIES, 0);
- acc0 = acc0.add(diff16.mul(diff16));
+ ShortVector prod16 = diff8.mul(diff8).reinterpretAsShorts();
+ acc0 = acc0.add(prod16.and((short) 0xFF).add(prod16.lanewise(LSHR, 8)));
// lower
ByteVector diff8a = vb8.lanewise(LSHR, 4).sub(va8.lanewise(LSHR, 4));
- Vector<Short> diff16a = diff8a.convertShape(B2S, Int4Constants.SHORT_SPECIES, 0);
- acc1 = acc1.add(diff16a.mul(diff16a));
+ ShortVector prod16a = diff8a.mul(diff8a).reinterpretAsShorts();
+ acc1 = acc1.add(prod16a.and((short) 0xFF).add(prod16a.lanewise(LSHR, 8)));
}
- Vector<Integer> intAcc0 = acc0.convert(S2I, 0);
- Vector<Integer> intAcc1 = acc0.convert(S2I, 1);
- Vector<Integer> intAcc2 = acc1.convert(S2I, 0);
- Vector<Integer> intAcc3 = acc1.convert(S2I, 1);
- sum += intAcc0.add(intAcc1).add(intAcc2).add(intAcc3).reinterpretAsInts().reduceLanes(ADD);
+
+ IntVector intAcc0 = acc0.reinterpretAsInts();
+ IntVector intAcc1 = acc1.reinterpretAsInts();
+ sum +=
+ intAcc0
+ .and(0xFFFF)
+ .add(intAcc0.lanewise(LSHR, 16))
+ .add(intAcc1.and(0xFFFF))
+ .add(intAcc1.lanewise(LSHR, 16))
+ .reduceLanes(ADD);
}
return sum;
}