add topk

Author: Joe Petviashvili

src/common.jl

@@ -23,13 +23,13 @@ export AF_NORM_MATRIX_1, AF_NORM_MATRIX_2, AF_NORM_MATRIX_INF, AF_NORM_MATRIX_L_
 export AF_NORM_VECTOR_INF, AF_NORM_VECTOR_P, AF_PAD_SYM, AF_PAD_ZERO, AF_RANDOM_ENGINE_DEFAULT, AF_RANDOM_ENGINE_MERSENNE
 export AF_RANDOM_ENGINE_MERSENNE_GP11213, AF_RANDOM_ENGINE_PHILOX, AF_RANDOM_ENGINE_PHILOX_4X32_10, AF_RANDOM_ENGINE_THREEFRY
 export AF_RANDOM_ENGINE_THREEFRY_2X32_16, AF_RGB, AF_SAD, AF_SHD, AF_SSD, AF_STORAGE_COO, AF_STORAGE_CSC
-export AF_STORAGE_CSR, AF_STORAGE_DENSE, AF_SUCCESS, AF_YCC_2020, AF_YCC_601, AF_YCC_709, AF_YCbCr, AF_ZNCC
-export AF_ZSAD, AF_ZSSD, afDevice, afHost, af_array, af_backend, af_binary_op, af_border_type, af_canny_threshold
-export af_colormap, af_connectivity, af_conv_domain, af_conv_mode, af_cspace_t, af_diffusion_eq, af_dtype
-export af_err, af_features, af_flux_function, af_homography_type, af_image_format, af_interp_type, af_marker_type
-export af_mat_prop, af_match_type, af_moment_type, af_norm_type, af_random_engine, af_random_engine_type
-export af_someenum_t, af_source, af_storage, af_window, af_ycc_std, b8, c32, c64, dim_t, f32, f64, intl
-export s16, s32, s64, u16, u32, u64, u8, uintl
+export AF_STORAGE_CSR, AF_STORAGE_DENSE, AF_SUCCESS, AF_TOPK_DEFAULT, AF_TOPK_MAX, AF_TOPK_MIN, AF_YCC_2020
+export AF_YCC_601, AF_YCC_709, AF_YCbCr, AF_ZNCC, AF_ZSAD, AF_ZSSD, afDevice, afHost, af_array, af_backend
+export af_binary_op, af_border_type, af_canny_threshold, af_colormap, af_connectivity, af_conv_domain
+export af_conv_mode, af_cspace_t, af_diffusion_eq, af_dtype, af_err, af_features, af_flux_function, af_homography_type
+export af_image_format, af_interp_type, af_marker_type, af_mat_prop, af_match_type, af_moment_type, af_norm_type
+export af_random_engine, af_random_engine_type, af_someenum_t, af_source, af_storage, af_topk_function
+export af_window, af_ycc_std, b8, c32, c64, dim_t, f32, f64, intl, s16, s32, s64, u16, u32, u64, u8, uintl
 
 const dim_t = Clonglong
 const intl = Clonglong
@@ -295,6 +295,13 @@ const AF_DIFFUSION_MCDE = (UInt32)(2)
 const AF_DIFFUSION_DEFAULT = (UInt32)(0)
 # end enum af_diffusion_eq
 
+# begin enum af_topk_function
+const af_topk_function = UInt32
+const AF_TOPK_MAX = (UInt32)(1)
+const AF_TOPK_MIN = (UInt32)(2)
+const AF_TOPK_DEFAULT = (UInt32)(0)
+# end enum af_topk_function
+
 type af_seq
     _begin::Cdouble
     _end::Cdouble

src/wrap.jl

@@ -36,7 +36,7 @@ export set_device, set_intersect, set_manual_eval_flag, set_mem_step_size, set_p
 export set_title, set_union, set_unique, set_visibility, shift, show, sift, sigmoid, signbit, sin, sinh
 export skew, sobel_operator, solve, solve_lu, sort_by_key, sparse_convert_to, sparse_get_col_idx, sparse_get_info
 export sparse_get_nnz, sparse_get_row_idx, sparse_get_storage, sparse_get_values, sqrt, stdev_all, sub
-export sum, sum_all, sum_nan, sum_nan_all, susan, svd_inplace, sync, tan, tanh, tgamma, tile, transform
+export sum, sum_all, sum_nan, sum_nan_all, susan, svd_inplace, sync, tan, tanh, tgamma, tile, topk, transform
 export transform_coordinates, translate, transpose_inplace, trunc, unlock_array, unlock_device_ptr, unwrap
 export upper, var_all, var_all_weighted, wrap, write_array, ycbcr2rgb
 
@@ -1850,6 +1850,13 @@ function corrcoef(X::AFArray,Y::AFArray)
     (realVal[],imagVal[])
 end
 
+function topk(_in::AFArray,k::Integer,dim::Integer,order::af_topk_function)
+    values = RefValue{af_array}(0)
+    indices = RefValue{af_array}(0)
+    _error(ccall((:af_topk,af_lib),af_err,(Ptr{af_array},Ptr{af_array},af_array,Cint,Cint,af_topk_function),values,indices,_in.arr,Cint(k),Cint(dim - 1),order))
+    (AFArray!(values[]),AFArray!(indices[]))
+end
+
 function fast(_in::AFArray,thr::Cfloat,arc_length::Integer,non_max::Bool,feature_ratio::Cfloat,edge::Integer)
     out = RefValue{af_features}(0)
     _error(ccall((:af_fast,af_lib),af_err,(Ptr{af_features},af_array,Cfloat,UInt32,Bool,Cfloat,UInt32),out,_in.arr,thr,UInt32(arc_length),non_max,feature_ratio,UInt32(edge)))