topk: fix STL compatibility, cstep indexing, omp barrier, and code style

- Guard <algorithm>/<vector> behind #if NCNN_SIMPLESTL, include simplestl.h
- Use std::partial_sort in simplestl mode (no std::nth_element available)
- Guard <math.h> in tests behind #if !NCNN_SIMPLESTL to avoid simplemath.h
  conflict; define INFINITY/NAN as float expressions in simplestl mode
- Fix cstep-unaware indexing for 3D/4D output tensors: use actual cstep
  for channel offset instead of assuming contiguous w*h layout
- Convert #pragma omp parallel + inner #pragma omp for to #pragma omp
  parallel for to avoid __kmpc_barrier in simpleomp mode
- Fix copyright year 2026->2025
- Apply code-format whitespace cleanup

Author: vlordier

src/layer/topk.cpp

@@ -1,12 +1,17 @@
-// Copyright 2026 Tencent
+// Copyright 2025 Tencent
 // SPDX-License-Identifier: BSD-3-Clause
 
 #include "topk.h"
 
-#include <algorithm>
 #include <stdint.h>
 #include <string.h>
+
+#if NCNN_SIMPLESTL
+#include "simplestl.h"
+#else
+#include <algorithm>
 #include <vector>
+#endif
 
 #if __ARM_NEON
 #include <arm_neon.h>
@@ -185,6 +190,21 @@ int TopK::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_bl
 
     const int total_lines = outer * inner;
 
+    // ncnn 3-/4-D mats have a channel stride (cstep) that may be larger than w*h
+    // due to alignment padding.  The flat inner/outer indexing must account for this:
+    //   - when axis reduces a non-channel dim, the outer loop spans channels and
+    //     the channel offset must use cstep rather than the product of spatial sizes;
+    //   - when axis IS the channel dim, the per-element j-stride must be cstep.
+    const size_t in_cstep = (dims >= 3) ? (size_t)bottom_blob.cstep : 0;
+    const size_t out_cstep = (dims >= 3) ? values.cstep : 0;
+    const bool axis_is_channel = (dims >= 3 && positive_axis == dims - 1);
+    // spatial-only outer count: channels factored out so cstep can be used separately
+    const int c_channels = (!axis_is_channel && dims >= 3) ? shape[dims - 1] : 1;
+    const int outer_spatial = (dims >= 3 && !axis_is_channel) ? outer / c_channels : outer;
+    // stride when stepping along the axis in memory
+    const size_t in_axis_stride = axis_is_channel ? in_cstep : (size_t)inner;
+    const size_t out_axis_stride = axis_is_channel ? out_cstep : (size_t)inner;
+
     if (_k == 1)
     {
         #pragma omp parallel for num_threads(opt.num_threads)
@@ -193,8 +213,19 @@ int TopK::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_bl
             int outer_i = line / inner;
             int inner_i = line - outer_i * inner;
 
-            int in_base = outer_i * axis_size * inner + inner_i;
-            int out_base = outer_i * inner + inner_i;
+            size_t in_base, out_base;
+            if (!axis_is_channel && dims >= 3)
+            {
+                const int ci = outer_i / outer_spatial;
+                const int sp_i = outer_i % outer_spatial;
+                in_base = (size_t)ci * in_cstep + (size_t)sp_i * axis_size * inner + inner_i;
+                out_base = (size_t)ci * out_cstep + (size_t)sp_i * 1 * inner + inner_i;
+            }
+            else
+            {
+                in_base = (size_t)outer_i * axis_size * inner + inner_i;
+                out_base = (size_t)outer_i * 1 * inner + inner_i;
+            }
 
 #if __ARM_NEON
             if (!output_indices && inner == 1 && axis_size >= 4)
@@ -273,7 +304,7 @@ int TopK::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_bl
 
             for (int j = 1; j < axis_size; j++)
             {
-                const float candidate_value = ptr[in_base + j * inner];
+                const float candidate_value = ptr[in_base + j * in_axis_stride];
                 if (topk_value_index_comp(candidate_value, j, best_value, best_index, largest_flag))
                 {
                     best_value = candidate_value;
@@ -301,22 +332,33 @@ int TopK::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_bl
             int outer_i = line / inner;
             int inner_i = line - outer_i * inner;
 
-            int in_base = outer_i * axis_size * inner + inner_i;
-            int out_base = outer_i * _k * inner + inner_i;
+            size_t in_base, out_base;
+            if (!axis_is_channel && dims >= 3)
+            {
+                const int ci = outer_i / outer_spatial;
+                const int sp_i = outer_i % outer_spatial;
+                in_base = (size_t)ci * in_cstep + (size_t)sp_i * axis_size * inner + inner_i;
+                out_base = (size_t)ci * out_cstep + (size_t)sp_i * _k * inner + inner_i;
+            }
+            else
+            {
+                in_base = (size_t)outer_i * axis_size * inner + inner_i;
+                out_base = (size_t)outer_i * _k * inner + inner_i;
+            }
 
             if (output_indices)
             {
                 for (int j = 0; j < _k; j++)
                 {
-                    outptr[out_base + j * inner] = ptr[in_base + j * inner];
-                    outidxptr[out_base + j * inner] = (float)j;
+                    outptr[out_base + j * out_axis_stride] = ptr[in_base + j * in_axis_stride];
+                    outidxptr[out_base + j * out_axis_stride] = (float)j;
                 }
             }
             else
             {
                 for (int j = 0; j < _k; j++)
                 {
-                    outptr[out_base + j * inner] = ptr[in_base + j * inner];
+                    outptr[out_base + j * out_axis_stride] = ptr[in_base + j * in_axis_stride];
                 }
             }
         }
@@ -336,8 +378,19 @@ int TopK::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_bl
             int outer_i = line / inner;
             int inner_i = line - outer_i * inner;
 
-            int in_base = outer_i * axis_size * inner + inner_i;
-            int out_base = outer_i * _k * inner + inner_i;
+            size_t in_base, out_base;
+            if (!axis_is_channel && dims >= 3)
+            {
+                const int ci = outer_i / outer_spatial;
+                const int sp_i = outer_i % outer_spatial;
+                in_base = (size_t)ci * in_cstep + (size_t)sp_i * axis_size * inner + inner_i;
+                out_base = (size_t)ci * out_cstep + (size_t)sp_i * _k * inner + inner_i;
+            }
+            else
+            {
+                in_base = (size_t)outer_i * axis_size * inner + inner_i;
+                out_base = (size_t)outer_i * _k * inner + inner_i;
+            }
 
             float top_values[4];
             int top_indices[4];
@@ -347,7 +400,7 @@ int TopK::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_bl
             {
                 for (int j = 0; j < axis_size; j++)
                 {
-                    const float candidate_value = ptr[in_base + j * inner];
+                    const float candidate_value = ptr[in_base + j * in_axis_stride];
 
                     if (top_count < _k)
                     {
@@ -382,7 +435,7 @@ int TopK::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_bl
             {
                 for (int j = 0; j < axis_size; j++)
                 {
-                    const float candidate_value = ptr[in_base + j * inner];
+                    const float candidate_value = ptr[in_base + j * in_axis_stride];
 
                     if (top_count < _k)
                     {
@@ -412,15 +465,15 @@ int TopK::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_bl
             {
                 for (int j = 0; j < _k; j++)
                 {
-                    outptr[out_base + j * inner] = top_values[j];
-                    outidxptr[out_base + j * inner] = (float)top_indices[j];
+                    outptr[out_base + j * out_axis_stride] = top_values[j];
+                    outidxptr[out_base + j * out_axis_stride] = (float)top_indices[j];
                 }
             }
             else
             {
                 for (int j = 0; j < _k; j++)
                 {
-                    outptr[out_base + j * inner] = top_values[j];
+                    outptr[out_base + j * out_axis_stride] = top_values[j];
                 }
             }
         }
@@ -432,58 +485,73 @@ int TopK::forward(const std::vector<Mat>& bottom_blobs, std::vector<Mat>& top_bl
         return 0;
     }
 
-    #pragma omp parallel num_threads(opt.num_threads)
+    #pragma omp parallel for num_threads(opt.num_threads)
+    for (int line = 0; line < total_lines; line++)
     {
-        std::vector<std::pair<float, int> > vec;
-        vec.resize(axis_size);
+        std::vector<std::pair<float, int> > vec(axis_size);
 
         topk_pair_comparator comp(largest_flag);
 
-        #pragma omp for
-        for (int line = 0; line < total_lines; line++)
-        {
-            int outer_i = line / inner;
-            int inner_i = line - outer_i * inner;
+        int outer_i = line / inner;
+        int inner_i = line - outer_i * inner;
 
-            int in_base = outer_i * axis_size * inner + inner_i;
-            int out_base = outer_i * _k * inner + inner_i;
+        size_t in_base, out_base;
+        if (!axis_is_channel && dims >= 3)
+        {
+            const int ci = outer_i / outer_spatial;
+            const int sp_i = outer_i % outer_spatial;
+            in_base = (size_t)ci * in_cstep + (size_t)sp_i * axis_size * inner + inner_i;
+            out_base = (size_t)ci * out_cstep + (size_t)sp_i * _k * inner + inner_i;
+        }
+        else
+        {
+            in_base = (size_t)outer_i * axis_size * inner + inner_i;
+            out_base = (size_t)outer_i * _k * inner + inner_i;
+        }
 
-            for (int j = 0; j < axis_size; j++)
-            {
-                vec[j].first = ptr[in_base + j * inner];
-                vec[j].second = j;
-            }
+        for (int j = 0; j < axis_size; j++)
+        {
+            vec[j].first = ptr[in_base + j * in_axis_stride];
+            vec[j].second = j;
+        }
 
-            if (_k < axis_size)
+        if (_k < axis_size)
+        {
+#if NCNN_SIMPLESTL
+            std::partial_sort(vec.begin(), vec.begin() + _k, vec.end(), comp);
+#else
+            if (sorted_flag)
             {
-                if (sorted_flag)
-                {
-                    std::nth_element(vec.begin(), vec.begin() + _k, vec.end(), comp);
-                    std::sort(vec.begin(), vec.begin() + _k, comp);
-                }
-                else
-                    std::nth_element(vec.begin(), vec.begin() + _k, vec.end(), comp);
+                std::nth_element(vec.begin(), vec.begin() + _k, vec.end(), comp);
+                std::sort(vec.begin(), vec.begin() + _k, comp);
             }
             else
-            {
-                if (sorted_flag)
-                    std::sort(vec.begin(), vec.end(), comp);
-            }
+                std::nth_element(vec.begin(), vec.begin() + _k, vec.end(), comp);
+#endif
+        }
+        else
+        {
+            if (sorted_flag)
+#if NCNN_SIMPLESTL
+                std::partial_sort(vec.begin(), vec.end(), vec.end(), comp);
+#else
+                std::sort(vec.begin(), vec.end(), comp);
+#endif
+        }
 
-            if (output_indices)
+        if (output_indices)
+        {
+            for (int j = 0; j < _k; j++)
             {
-                for (int j = 0; j < _k; j++)
-                {
-                    outptr[out_base + j * inner] = vec[j].first;
-                    outidxptr[out_base + j * inner] = (float)vec[j].second;
-                }
+                outptr[out_base + j * out_axis_stride] = vec[j].first;
+                outidxptr[out_base + j * out_axis_stride] = (float)vec[j].second;
             }
-            else
+        }
+        else
+        {
+            for (int j = 0; j < _k; j++)
             {
-                for (int j = 0; j < _k; j++)
-                {
-                    outptr[out_base + j * inner] = vec[j].first;
-                }
+                outptr[out_base + j * out_axis_stride] = vec[j].first;
             }
         }
     }

src/layer/topk.h

@@ -1,4 +1,4 @@
-// Copyright 2026 Tencent
+// Copyright 2025 Tencent
 // SPDX-License-Identifier: BSD-3-Clause
 
 #ifndef LAYER_TOPK_H

tests/test_topk.cpp

@@ -1,9 +1,17 @@
-// Copyright 2026 Tencent
+// Copyright 2025 Tencent
 // SPDX-License-Identifier: BSD-3-Clause
 
 #include "testutil.h"
 
-#include <limits>
+#if NCNN_SIMPLESTL
+// simplemath.h conflicts with system math.h; define only what we need
+static const float TEST_INF = 1.f / 0.f;
+static const float TEST_NAN = 0.f / 0.f;
+#define INFINITY TEST_INF
+#define NAN      TEST_NAN
+#else
+#include <math.h>
+#endif
 
 static int test_topk_cpu_forward(const ncnn::Mat& a, int axis, int k, int largest, int sorted, ncnn::Mat& values, ncnn::Mat& indices)
 {
@@ -121,7 +129,7 @@ static int test_topk_0()
     return 0
            || test_topk(a, 0, 1, 1, 1)
            || test_topk(a, 0, 5, 1, 1)
-            || test_topk(a, 0, 1, 0, 0)
+           || test_topk(a, 0, 1, 0, 0)
            || test_topk(a, -1, 7, 0, 1)
            || test_topk(a, 0, 4, 1, 0)
            || test_topk(a, 0, 9, 1, 1);
@@ -175,9 +183,9 @@ static int test_topk_inf_order()
     ncnn::Mat a(6);
     float* ptr = a;
     ptr[0] = 1.f;
-    ptr[1] = std::numeric_limits<float>::infinity();
+    ptr[1] = INFINITY;
     ptr[2] = -2.f;
-    ptr[3] = -std::numeric_limits<float>::infinity();
+    ptr[3] = -INFINITY;
     ptr[4] = 0.5f;
     ptr[5] = 3.f;
 
@@ -193,7 +201,7 @@ static int test_topk_inf_order()
 
     const float* vptr = values;
     const float* iptr = indices;
-    if (values.w != 2 || indices.w != 2 || vptr[0] != std::numeric_limits<float>::infinity() || vptr[1] != 3.f || (int)iptr[0] != 1 || (int)iptr[1] != 5)
+    if (values.w != 2 || indices.w != 2 || vptr[0] != INFINITY || vptr[1] != 3.f || (int)iptr[0] != 1 || (int)iptr[1] != 5)
     {
         fprintf(stderr, "test_topk_inf_order largest result mismatch\n");
         return -1;
@@ -208,7 +216,7 @@ static int test_topk_inf_order()
 
     vptr = values;
     iptr = indices;
-    if (values.w != 2 || indices.w != 2 || vptr[0] != -std::numeric_limits<float>::infinity() || vptr[1] != -2.f || (int)iptr[0] != 3 || (int)iptr[1] != 2)
+    if (values.w != 2 || indices.w != 2 || vptr[0] != -INFINITY || vptr[1] != -2.f || (int)iptr[0] != 3 || (int)iptr[1] != 2)
     {
         fprintf(stderr, "test_topk_inf_order smallest result mismatch\n");
         return -1;
@@ -222,7 +230,7 @@ static int test_topk_nan_robust()
     ncnn::Mat a(4);
     float* ptr = a;
     ptr[0] = 1.f;
-    ptr[1] = std::numeric_limits<float>::quiet_NaN();
+    ptr[1] = NAN;
     ptr[2] = 2.f;
     ptr[3] = -1.f;
 

tools/pnnx/src/ir.cpp

@@ -1640,12 +1640,12 @@ int Graph::python(const std::string& pypath, const std::string& pnnxbinpath, con
                 continue;
 
             fprintf(pyfp, "        self.%s = TopK(", sanitize_identifier(op->name).c_str());
-            
+
             int i = 0;
             for (const auto& it : op->params)
             {
                 fprintf(pyfp, "%s=", it.first.c_str());
-                
+
                 const Parameter& param = it.second;
                 if (param.type == 2)
                 {
@@ -1655,12 +1655,12 @@ int Graph::python(const std::string& pypath, const std::string& pnnxbinpath, con
                 {
                     fprintf(pyfp, "%d", param.b ? 1 : 0);
                 }
-                
+
                 if (i + 1 != op->params.size())
                     fprintf(pyfp, ", ");
                 i++;
             }
-            
+
             fprintf(pyfp, ")\n");
         }
     }

tools/pnnx/src/pass_ncnn/TopK.cpp

@@ -1,4 +1,4 @@
-// Copyright 2026 Tencent
+// Copyright 2025 Tencent
 // SPDX-License-Identifier: BSD-3-Clause
 
 #include "pass_ncnn.h"