Vectorize argmax with TensorPrimitives

Directory.Packages.props

@@ -8,6 +8,7 @@
     <PackageVersion Include="Microsoft.ML.Tokenizers" Version="2.0.0" />
     <PackageVersion Include="Microsoft.ML.Tokenizers.Data.O200kBase" Version="2.0.0" />
     <PackageVersion Include="Microsoft.Bcl.Memory" Version="10.0.10" />
+    <PackageVersion Include="System.Numerics.Tensors" Version="10.0.10" />
   </ItemGroup>
 
   <ItemGroup Label="Tests">

bench/PrivacyFilter.Net.Benchmarks/Program.cs

@@ -97,7 +97,39 @@ public int[] ViterbiDense() =>
     public int[] ViterbiSparse() => _decoder.Decode(_logProbabilities, TokenCount);
 
     [Benchmark]
-    public int[] ArgMax() => PrivacyFilter.ArgMax(_logProbabilities, TokenCount, _classCount);
+    public int[] ArgMaxScalar() =>
+        ArgMaxScalar(_logProbabilities, TokenCount, _classCount);
+
+    [Benchmark]
+    public int[] ArgMaxTensorPrimitives() =>
+        PrivacyFilter.ArgMax(_logProbabilities, TokenCount, _classCount);
+
+    private static int[] ArgMaxScalar(
+        ReadOnlySpan<float> scores,
+        int tokenCount,
+        int classCount)
+    {
+        var labels = new int[tokenCount];
+        for (int token = 0; token < tokenCount; token++)
+        {
+            int offset = token * classCount;
+            int bestLabel = 0;
+            float bestScore = scores[offset];
+            for (int label = 1; label < classCount; label++)
+            {
+                float score = scores[offset + label];
+                if (score > bestScore)
+                {
+                    bestScore = score;
+                    bestLabel = label;
+                }
+            }
+
+            labels[token] = bestLabel;
+        }
+
+        return labels;
+    }
 
     private static int[] DecodeDense(
         float[] emissions,

src/PrivacyFilter.Net/PrivacyFilter.Net.csproj

@@ -47,6 +47,7 @@
     <PackageReference Include="Microsoft.ML.Tokenizers" />
     <PackageReference Include="Microsoft.ML.Tokenizers.Data.O200kBase" />
     <PackageReference Include="Microsoft.Bcl.Memory" />
+    <PackageReference Include="System.Numerics.Tensors" />
   </ItemGroup>
 
 </Project>

src/PrivacyFilter.Net/PrivacyFilter.cs

@@ -1,6 +1,7 @@
 using System.Buffers;
 using Microsoft.ML.OnnxRuntime;
 using Microsoft.ML.OnnxRuntime.Tensors;
+using TensorPrimitives = System.Numerics.Tensors.TensorPrimitives;
 
 namespace PrivacyFilterNet;
 
@@ -247,20 +248,8 @@ internal static int[] ArgMax(ReadOnlySpan<float> scores, int tokenCount, int cla
         var labels = new int[tokenCount];
         for (int token = 0; token < tokenCount; token++)
         {
-            int offset = token * classCount;
-            int bestLabel = 0;
-            float bestScore = scores[offset];
-            for (int label = 1; label < classCount; label++)
-            {
-                float score = scores[offset + label];
-                if (score > bestScore)
-                {
-                    bestScore = score;
-                    bestLabel = label;
-                }
-            }
-
-            labels[token] = bestLabel;
+            labels[token] = TensorPrimitives.IndexOfMax(
+                scores.Slice(token * classCount, classCount));
         }
 
         return labels;

tests/PrivacyFilter.Net.Tests/DecoderTests.cs

@@ -84,6 +84,28 @@ public void SparseViterbiMatchesDenseReference()
             decoder.Decode(tiedEmissions, 8));
     }
 
+    [Fact]
+    public void TensorArgMaxMatchesScalarForFiniteScores()
+    {
+        const int tokenCount = 128;
+        int classCount = ClassNames.Length;
+        var scores = new float[tokenCount * classCount];
+        var random = new Random(42);
+        for (int index = 0; index < scores.Length; index++)
+        {
+            scores[index] = (random.NextSingle() * 20) - 10;
+        }
+
+        Assert.Equal(
+            ArgMaxScalar(scores, tokenCount, classCount),
+            PrivacyFilter.ArgMax(scores, tokenCount, classCount));
+
+        Array.Clear(scores);
+        Assert.Equal(
+            ArgMaxScalar(scores, tokenCount, classCount),
+            PrivacyFilter.ArgMax(scores, tokenCount, classCount));
+    }
+
     [Fact]
     public void SpanDecoderTrimsAndRedacts()
     {
@@ -113,6 +135,33 @@ public void SpanDecoderTrimsAndRedacts()
 
     private static int ClassIndex(string name) => Array.IndexOf(ClassNames, name);
 
+    private static int[] ArgMaxScalar(
+        ReadOnlySpan<float> scores,
+        int tokenCount,
+        int classCount)
+    {
+        var labels = new int[tokenCount];
+        for (int token = 0; token < tokenCount; token++)
+        {
+            int offset = token * classCount;
+            int bestLabel = 0;
+            float bestScore = scores[offset];
+            for (int label = 1; label < classCount; label++)
+            {
+                float score = scores[offset + label];
+                if (score > bestScore)
+                {
+                    bestScore = score;
+                    bestLabel = label;
+                }
+            }
+
+            labels[token] = bestLabel;
+        }
+
+        return labels;
+    }
+
     private static int[] DecodeDenseReference(float[] emissions, int tokenCount, int classCount)
     {
         const float invalid = -1e9f;