1 files changed, 98 insertions, 0 deletions
diff --git a/runtimes/libs/ARMComputeEx/src/core/CL/cl_kernels/topkv2.cl b/runtimes/libs/ARMComputeEx/src/core/CL/cl_kernels/topkv2.cl
new file mode 100644
index 000000000..50472e4f9
--- /dev/null
+++ b/runtimes/libs/ARMComputeEx/src/core/CL/cl_kernels/topkv2.cl
@@ -0,0 +1,98 @@
+/*
+ * Copyright (c) 2018 Samsung Electronics Co., Ltd. All Rights Reserved
+ * Copyright (c) 2017 ARM Limited.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "helpers.h"
+
+__kernel void topkv2_init(VECTOR_DECLARATION(input), __global float *in_key_buf,
+                          __global int *in_ind_buf, const int n)
+{
+  int gid = get_global_id(0);
+  int lws = get_local_size(0);
+  int groups = get_num_groups(0);
+  int gws = lws * groups;
+  int iter = n / gws;
+
+  Vector input = CONVERT_TO_VECTOR_STRUCT_NO_STEP(input);
+
+  for (int i = 0; i < iter; ++i)
+  {
+    int idx = i * gws + gid;
+    in_key_buf[idx] = *(__global float *)(input.ptr + idx * input.stride_x);
+    in_ind_buf[idx] = idx;
+  }
+}
+
+__kernel void topkv2_find_first_negative(__global float *out_key_buf,
+                                         __global int *first_negative_idx, int n)
+{
+  int gid = get_global_id(0);
+
+  if (gid == n - 1)
+  {
+    // if the last item is positive, the first negative index is n.
+    if (out_key_buf[gid] > 0.f)
+      *first_negative_idx = n;
+  }
+  else if (gid == 0)
+  {
+    // if the first item is negative, set it 0.
+    if (out_key_buf[gid] < 0.f)
+      *first_negative_idx = 0;
+  }
+  else
+  {
+    // if its left is positive and it is negative, then it is the first negative item.
+    if (out_key_buf[gid - 1] > 0.f && out_key_buf[gid] < 0.f)
+      *first_negative_idx = gid;
+  }
+}
+
+__kernel void topkv2_reorder_negatives(__global float *in_key_buf, __global float *out_key_buf,
+                                       __global float *in_ind_buf, __global float *out_ind_buf,
+                                       __global int *first_negative_idx, int n)
+{
+  int gid = get_global_id(0);
+
+  int num_negs = n - *first_negative_idx;
+  int in_idx;
+
+  if (gid < num_negs)
+  {
+    in_idx = n - 1 - gid;
+  }
+  else
+  {
+    in_idx = gid - num_negs;
+  }
+
+  out_key_buf[gid] = in_key_buf[in_idx];
+  out_ind_buf[gid] = in_ind_buf[in_idx];
+}
+
+__kernel void topkv2_store(VECTOR_DECLARATION(values), VECTOR_DECLARATION(indices),
+                           __global float *out_key_buf, __global int *out_ind_buf, int n)
+{
+  int gid = get_global_id(0);
+
+  Vector values = CONVERT_TO_VECTOR_STRUCT_NO_STEP(values);
+  Vector indices = CONVERT_TO_VECTOR_STRUCT_NO_STEP(indices);
+
+  int idx = n - 1 - gid;
+
+  *(__global float *)(values.ptr + gid * values.stride_x) = out_key_buf[idx];
+  *(__global int *)(indices.ptr + gid * indices.stride_x) = out_ind_buf[idx];
+}