Update gather jit op

Pair programming with Paul

src/targets/gpu/jit/gather.cpp

@@ -47,8 +47,7 @@ extern "C" {
 __global__ void gather_kernel(void* in_data, void* in_indices, void* output) 
 {
     make_tensors()(in_data, in_indices, output)([](auto&&... xs) { 
-        auto settings = make_gather_settings(MIGRAPHX_MAKE_CONSTANT(int64_t{AXIS}));
-        gather(xs..., settings); 
+        gather<${axis}>(xs...); 
     });
 }
 
@@ -58,7 +57,7 @@ __global__ void gather_kernel(void* in_data, void* in_indices, void* output)
 
 )__migraphx__";
 
-struct gathernd_compiler : compiler<gathernd_compiler>
+struct gather_compiler : compiler<gather_compiler>
 {
     std::vector<std::string> names() const { return {"gather"}; }
 
@@ -72,12 +71,13 @@ struct gathernd_compiler : compiler<gathernd_compiler>
         options.kernel_name    = "gather_kernel";
         options.virtual_inputs = inputs;
 
-        // axis attribute
         assert(v.contains("axis"));
-        auto axis = v.at("axis").to<int64_t>();
-        options.params += " -DAXIS=" + std::to_string(axis);
+        auto axis = v.at("axis").to<std::string>();
+
+        auto src = interpolate_string(gather_kernel, {{"axis", axis}});
 
-        return compile_hip_code_object(gather_kernel, options);
+        // axis attribute
+        return compile_hip_code_object(src, options);
     }
 
     compiler_replace compile(context& ctx, instruction_ref ins, const operation& op) const

src/targets/gpu/kernels/include/migraphx/kernels/gather.hpp

@@ -29,41 +29,19 @@
 
 namespace migraphx {
 
-template <class T>
-struct gather_settings
-{
-    T axes{};
-};
-
-template <class... Ts>
-constexpr gather_settings<Ts...> make_gather_settings(Ts... xs)
-{
-    return {xs...};
-}
-
-template <class T, class U, class V, class Settings>
-__device__ void gather(const T& data_t, const U& indices_t, const V& output_t, Settings s)
+template <int axis, class T, class U, class V>
+__device__ void gather(const T& data_t, const U& indices_t, const V& output_t)
 {
     auto ind           = make_index();
-    auto axis          = s.axis;
     auto output_shape  = output_t.get_shape();
-    auto indices_shape = indices_t.get_shape();
-    auto data_shape    = data_t.get_shape();
-
-    auto axis_dim_size = data_shape.lens().at(axis);
-
-    auto indices_shape_lens = indices_shape.lens;
-    auto data_shape_lens    = data_shape.lens;
-
-    const auto* indices_ptr = indices_t.data();
-    auto* output_ptr        = output_t.data();
+    auto axis_dim_size = data_t.get_shape().lens[axis];
 
     ind.global_stride(output_shape.elements(), [&](auto i) {
         auto idx      = output_shape.multi(i);
-        auto in_index = indices_ptr[idx[axis]];
+        auto in_index = indices_t[idx[axis]];
         in_index      = (in_index < 0) ? in_index + axis_dim_size : in_index;
         idx[axis]     = in_index;
-        output_ptr[i] = indices_t[idx];
+        output_t[i]   = indices_t[idx];
     });
 }