merge changes from develop branch

src/targets/gpu/hip.cpp

@@ -109,10 +109,9 @@ argument register_on_gpu(const argument& arg)
 {
     auto arg_shared = arg.share();
     auto p          = share(register_on_gpu(arg_shared.data(), arg_shared.get_shape().bytes()));
-    return {arg_shared.get_shape(),
-            [ p, a = std::move(arg_shared) ]() mutable {return get_device_ptr(p.get());
-}
-}; // namespace gpu
+    return {arg_shared.get_shape(), [p, a = std::move(arg_shared)]() mutable {
+                return get_device_ptr(p.get());
+            }}; // namespace gpu
 } // namespace MIGRAPHX_INLINE_NS
 
 argument to_gpu(const argument& arg, bool host)

src/targets/gpu/include/migraphx/gpu/compile_scatternd.hpp

+#ifndef MIGRAPHX_GUARD_GPU_COMPILE_SCATTERND_HPP
+#define MIGRAPHX_GUARD_GPU_COMPILE_SCATTERND_HPP
+
+#include <migraphx/config.hpp>
+#include <migraphx/operation.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+
+struct context;
+operation
+compile_scatternd(context& ctx, const std::vector<shape>& io_shapes, const std::string& reduction);
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+#endif // MIGRAPHX_GUARD_GPU_COMPILE_SCATTERND_HPP

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

@@ -21,6 +21,16 @@ struct greater
     }
 };
 
+template <class InputIt, class OutputIt>
+constexpr OutputIt copy(InputIt first, InputIt last, OutputIt d_first)
+{
+    while(first != last)
+    {
+        *d_first++ = *first++;
+    }
+    return d_first;
+}
+
 template <class Iterator, class Compare>
 constexpr Iterator is_sorted_until(Iterator first, Iterator last, Compare comp)
 {

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

@@ -48,7 +48,7 @@ MIGRAPHX_INTEGRAL_CONSTANT_BINARY_OP(>=)
 MIGRAPHX_INTEGRAL_CONSTANT_BINARY_OP(==)
 MIGRAPHX_INTEGRAL_CONSTANT_BINARY_OP(!=)
 MIGRAPHX_INTEGRAL_CONSTANT_BINARY_OP(&)
-MIGRAPHX_INTEGRAL_CONSTANT_BINARY_OP (^)
+MIGRAPHX_INTEGRAL_CONSTANT_BINARY_OP(^)
 MIGRAPHX_INTEGRAL_CONSTANT_BINARY_OP(|)
 MIGRAPHX_INTEGRAL_CONSTANT_BINARY_OP(&&)
 MIGRAPHX_INTEGRAL_CONSTANT_BINARY_OP(||)

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

@@ -59,6 +59,7 @@ MIGRAPHX_DEVICE_MATH(cosh, ::cosh)
 MIGRAPHX_DEVICE_MATH(erf, ::erf)
 MIGRAPHX_DEVICE_MATH(exp, ::exp)
 MIGRAPHX_DEVICE_MATH(floor, ::floor)
+MIGRAPHX_DEVICE_MATH(isnan, ::isnan)
 MIGRAPHX_DEVICE_MATH(log, ::log)
 MIGRAPHX_DEVICE_MATH(pow, ::pow)
 MIGRAPHX_DEVICE_MATH(round, ::round)
@@ -103,6 +104,7 @@ MIGRAPHX_DEVICE_MATH_HALF(cos, ::cos)
 MIGRAPHX_DEVICE_MATH_HALF(cosh, ::cosh)
 MIGRAPHX_DEVICE_MATH_HALF(erf, ::erf)
 MIGRAPHX_DEVICE_MATH_HALF(floor, ::floor)
+MIGRAPHX_DEVICE_MATH_HALF(isnan, ::isnan)
 MIGRAPHX_DEVICE_MATH_HALF(pow, ::pow)
 MIGRAPHX_DEVICE_MATH_HALF(round, ::round)
 MIGRAPHX_DEVICE_MATH_HALF(sin, ::sin)
@@ -129,6 +131,7 @@ MIGRAPHX_DEVICE_MATH_VEC(cosh)
 MIGRAPHX_DEVICE_MATH_VEC(erf)
 MIGRAPHX_DEVICE_MATH_VEC(exp)
 MIGRAPHX_DEVICE_MATH_VEC(floor)
+MIGRAPHX_DEVICE_MATH_VEC(isnan)
 MIGRAPHX_DEVICE_MATH_VEC(log)
 MIGRAPHX_DEVICE_MATH_VEC(pow)
 MIGRAPHX_DEVICE_MATH_VEC(round)

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

+#ifndef MIGRAPHX_GUARD_KERNELS_SCATTERND_HPP
+#define MIGRAPHX_GUARD_KERNELS_SCATTERND_HPP
+
+#include <migraphx/kernels/index.hpp>
+#include <migraphx/kernels/algorithm.hpp>
+
+namespace migraphx {
+
+struct assign_none
+{
+    template <class T, class U>
+    MIGRAPHX_DEVICE_CONSTEXPR void operator()(T& x, U y) const
+    {
+        x = y;
+    }
+};
+
+struct assign_add
+{
+    template <class T, class U>
+    MIGRAPHX_DEVICE_CONSTEXPR void operator()(T& x, U y) const
+    {
+        x += y;
+    }
+};
+
+struct assign_mul
+{
+    template <class T, class U>
+    MIGRAPHX_DEVICE_CONSTEXPR void operator()(T& x, U y) const
+    {
+        x *= y;
+    }
+};
+
+template <class T, class U, class V, class F>
+__device__ void scatternd(const T& indices_t, const U& updates_t, const V& output_t, F f)
+{
+    auto index         = make_index();
+    auto updates_shape = updates_t.get_shape();
+
+    index.global_stride(updates_shape.elements(), [&](auto i) {
+        auto output_shape = output_t.get_shape();
+
+        auto indices_shape = indices_t.get_shape();
+        auto k             = indices_shape.lens.back();
+        auto q             = indices_shape.lens.size();
+
+        auto updates_idx = updates_shape.multi(i);
+        auto indices_idx = indices_shape.multi(0);
+        copy(updates_idx.begin(), updates_idx.begin() + q - 1, indices_idx.begin());
+
+        auto index_start = indices_t.begin() + indices_shape.index(indices_idx);
+        auto index_end   = index_start + k;
+        auto out_idx     = output_shape.multi(0);
+        copy(index_start, index_end, out_idx.begin());
+        copy(updates_idx.begin() + q - 1, updates_idx.end(), out_idx.begin() + k);
+
+        f(output_t[out_idx], updates_t[i]);
+    });
+}
+
+} // namespace migraphx
+#endif

src/targets/gpu/lowering.cpp

@@ -21,6 +21,7 @@
 #include <migraphx/gpu/abs.hpp>
 #include <migraphx/gpu/batch_norm_inference.hpp>
 #include <migraphx/gpu/compile_roialign.hpp>
+#include <migraphx/gpu/compile_scatternd.hpp>
 #include <migraphx/gpu/context.hpp>
 #include <migraphx/gpu/convolution.hpp>
 #include <migraphx/gpu/deconvolution.hpp>
@@ -207,6 +208,7 @@ struct miopen_apply
         add_nms_op();
         add_quant_convolution_op();
         add_roialign();
+        add_scatternd();
     }
 
     void copy_params()
@@ -443,7 +445,6 @@ struct miopen_apply
                                             reshapes[2],
                                             reshapes[3],
                                             output);
-
         });
     }
 
@@ -503,7 +504,6 @@ struct miopen_apply
     void add_roialign()
     {
         apply_map.emplace("roialign", [=](instruction_ref ins) {
-
             auto s      = ins->get_shape();
             auto op_val = ins->get_operator().to_value();
             auto output = insert_allocation(ins, s);
@@ -516,6 +516,60 @@ struct miopen_apply
         });
     }
 
+    void add_scatternd()
+    {
+        apply_map.emplace("scatternd_none", [=](instruction_ref ins) {
+            auto s      = ins->get_shape();
+            auto op_val = ins->get_operator().to_value();
+            auto output = insert_allocation(ins, s);
+            auto args   = ins->inputs();
+            args.push_back(output);
+
+            auto io_shapes = to_shapes(args);
+            io_shapes.erase(io_shapes.begin());
+            const std::string reduction = "none";
+            auto co                     = compile_scatternd(get_context(), io_shapes, reduction);
+            auto copy   = mod->insert_instruction(ins, make_op("hip::copy"), args.front(), output);
+            args.back() = copy;
+            args.erase(args.begin());
+            return mod->replace_instruction(ins, co, args);
+        });
+
+        apply_map.emplace("scatternd_add", [=](instruction_ref ins) {
+            auto s      = ins->get_shape();
+            auto op_val = ins->get_operator().to_value();
+            auto output = insert_allocation(ins, s);
+            auto args   = ins->inputs();
+            args.push_back(output);
+
+            auto io_shapes = to_shapes(args);
+            io_shapes.erase(io_shapes.begin());
+            const std::string reduction = "add";
+            auto co                     = compile_scatternd(get_context(), io_shapes, reduction);
+            auto copy   = mod->insert_instruction(ins, make_op("hip::copy"), args.front(), output);
+            args.back() = copy;
+            args.erase(args.begin());
+            return mod->replace_instruction(ins, co, args);
+        });
+
+        apply_map.emplace("scatternd_mul", [=](instruction_ref ins) {
+            auto s      = ins->get_shape();
+            auto op_val = ins->get_operator().to_value();
+            auto output = insert_allocation(ins, s);
+            auto args   = ins->inputs();
+            args.push_back(output);
+
+            auto io_shapes = to_shapes(args);
+            io_shapes.erase(io_shapes.begin());
+            const std::string reduction = "mul";
+            auto co                     = compile_scatternd(get_context(), io_shapes, reduction);
+            auto copy   = mod->insert_instruction(ins, make_op("hip::copy"), args.front(), output);
+            args.back() = copy;
+            args.erase(args.begin());
+            return mod->replace_instruction(ins, co, args);
+        });
+    }
+
     // replace the loop operator with gpu_loop operator
     void add_loop_op()
     {

src/tf/tf_parser.cpp

@@ -499,8 +499,7 @@ literal tf_parser::parse_tensor(const tensorflow::TensorProto& t) const
         return create_literal(shape::int64_type, dims, get_data_vals(t.int64_val(), shape_size));
     case tensorflow::DataType::DT_BOOL:
         return create_literal(shape::int32_type, dims, get_data_vals(t.bool_val(), shape_size));
-    case tensorflow::DataType::DT_HALF:
-    {
+    case tensorflow::DataType::DT_HALF: {
         std::vector<int> data_int32 = get_data_vals(t.half_val(), shape_size);
         std::vector<uint16_t> data_uint16(data_int32.begin(), data_int32.end());
         std::vector<half> data_half;

test/include/test.hpp

@@ -197,7 +197,7 @@ struct lhs_expression
     TEST_LHS_REOPERATOR(%)
     TEST_LHS_REOPERATOR(&)
     TEST_LHS_REOPERATOR(|)
-    TEST_LHS_REOPERATOR (^)
+    TEST_LHS_REOPERATOR(^)
 };
 
 template <class F>

test/onnx/gen_onnx.py

@@ -2307,6 +2307,32 @@ def instance_norm_val_3d_test():
     return ([node], [], [y], [x_tensor, scale_tensor, bias_tensor])
 
 
+@onnx_test
+def isnan_float_test():
+    t1 = helper.make_tensor_value_info('t1', TensorProto.FLOAT, [2, 3])
+    t2 = helper.make_tensor_value_info('t2', TensorProto.FLOAT, [2, 3])
+
+    node = onnx.helper.make_node(
+        'IsNaN',
+        inputs=['t1'],
+        outputs=['t2'],
+    )
+    return ([node], [t1], [t2])
+
+
+@onnx_test
+def isnan_half_test():
+    t1 = helper.make_tensor_value_info('t1', TensorProto.FLOAT16, [2, 3])
+    t2 = helper.make_tensor_value_info('t2', TensorProto.FLOAT16, [2, 3])
+
+    node = onnx.helper.make_node(
+        'IsNaN',
+        inputs=['t1'],
+        outputs=['t2'],
+    )
+    return ([node], [t1], [t2])
+
+
 @onnx_test
 def layernorm_test():
     x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [1, 1, 5])
@@ -4126,6 +4152,59 @@ def scatter_test():
     return ([node], [x, i, u], [y])
 
 
+@onnx_test
+def scatternd_add_test():
+    data = helper.make_tensor_value_info('data', TensorProto.FLOAT, [2, 2, 2])
+    indices = helper.make_tensor_value_info('indices', TensorProto.INT64,
+                                            [2, 1, 2])
+    updates = helper.make_tensor_value_info('updates', TensorProto.FLOAT,
+                                            [2, 1, 2])
+    output = helper.make_tensor_value_info('output', TensorProto.FLOAT,
+                                           [2, 2, 2])
+
+    node = onnx.helper.make_node('ScatterND',
+                                 inputs=['data', 'indices', 'updates'],
+                                 outputs=['output'],
+                                 reduction="add")
+
+    return ([node], [data, indices, updates], [output])
+
+
+@onnx_test
+def scatternd_mul_test():
+    data = helper.make_tensor_value_info('data', TensorProto.FLOAT, [2, 2, 2])
+    indices = helper.make_tensor_value_info('indices', TensorProto.INT64,
+                                            [2, 1, 2])
+    updates = helper.make_tensor_value_info('updates', TensorProto.FLOAT,
+                                            [2, 1, 2])
+    output = helper.make_tensor_value_info('output', TensorProto.FLOAT,
+                                           [2, 2, 2])
+
+    node = onnx.helper.make_node('ScatterND',
+                                 inputs=['data', 'indices', 'updates'],
+                                 outputs=['output'],
+                                 reduction="mul")
+
+    return ([node], [data, indices, updates], [output])
+
+
+@onnx_test
+def scatternd_test():
+    data = helper.make_tensor_value_info('data', TensorProto.FLOAT, [2, 2, 2])
+    indices = helper.make_tensor_value_info('indices', TensorProto.INT64,
+                                            [2, 1, 2])
+    updates = helper.make_tensor_value_info('updates', TensorProto.FLOAT,
+                                            [2, 1, 2])
+    output = helper.make_tensor_value_info('output', TensorProto.FLOAT,
+                                           [2, 2, 2])
+
+    node = onnx.helper.make_node('ScatterND',
+                                 inputs=['data', 'indices', 'updates'],
+                                 outputs=['output'])
+
+    return ([node], [data, indices, updates], [output])
+
+
 @onnx_test
 def selu_test():
     x = helper.make_tensor_value_info('x', TensorProto.DOUBLE, [2, 3])

test/onnx/isnan_float_test.onnx

+isnan_float_test:O
+
+t1t2"IsNaNisnan_float_testZ
+t1
+

+
+b
+t2
+

+
+B
\ No newline at end of file

test/onnx/isnan_half_test.onnx

+isnan_half_test:N
+
+t1t2"IsNaNisnan_half_testZ
+t1
+
+
+
+b
+t2
+
+
+
+B
\ No newline at end of file

test/onnx/onnx_test.cpp

@@ -1929,6 +1929,32 @@ TEST_CASE(if_tuple_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(isnan_float_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape s{migraphx::shape::float_type, {2, 3}};
+    auto t1  = mm->add_parameter("t1", s);
+    auto ret = mm->add_instruction(migraphx::make_op("isnan"), t1);
+    mm->add_return({ret});
+
+    auto prog = migraphx::parse_onnx("isnan_float_test.onnx");
+    EXPECT(p == prog);
+}
+
+TEST_CASE(isnan_half_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape s{migraphx::shape::half_type, {2, 3}};
+    auto t1  = mm->add_parameter("t1", s);
+    auto ret = mm->add_instruction(migraphx::make_op("isnan"), t1);
+    mm->add_return({ret});
+
+    auto prog = migraphx::parse_onnx("isnan_half_test.onnx");
+    EXPECT(p == prog);
+}
+
 TEST_CASE(imagescaler_test)
 {
     migraphx::program p;
@@ -3970,6 +3996,57 @@ TEST_CASE(scatter_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(scatternd_test)
+{
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+        auto l0 =
+            mm->add_parameter("data", migraphx::shape{migraphx::shape::float_type, {2, 2, 2}});
+        auto l1 =
+            mm->add_parameter("indices", migraphx::shape{migraphx::shape::int64_type, {2, 1, 2}});
+        auto l2 =
+            mm->add_parameter("updates", migraphx::shape{migraphx::shape::float_type, {2, 1, 2}});
+        auto r = mm->add_instruction(migraphx::make_op("scatternd_none"), l0, l1, l2);
+        mm->add_return({r});
+        auto prog = migraphx::parse_onnx("scatternd_test.onnx");
+
+        EXPECT(p == prog);
+    }
+
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+        auto l0 =
+            mm->add_parameter("data", migraphx::shape{migraphx::shape::float_type, {2, 2, 2}});
+        auto l1 =
+            mm->add_parameter("indices", migraphx::shape{migraphx::shape::int64_type, {2, 1, 2}});
+        auto l2 =
+            mm->add_parameter("updates", migraphx::shape{migraphx::shape::float_type, {2, 1, 2}});
+        auto r = mm->add_instruction(migraphx::make_op("scatternd_add"), l0, l1, l2);
+        mm->add_return({r});
+        auto prog = migraphx::parse_onnx("scatternd_add_test.onnx");
+
+        EXPECT(p == prog);
+    }
+
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+        auto l0 =
+            mm->add_parameter("data", migraphx::shape{migraphx::shape::float_type, {2, 2, 2}});
+        auto l1 =
+            mm->add_parameter("indices", migraphx::shape{migraphx::shape::int64_type, {2, 1, 2}});
+        auto l2 =
+            mm->add_parameter("updates", migraphx::shape{migraphx::shape::float_type, {2, 1, 2}});
+        auto r = mm->add_instruction(migraphx::make_op("scatternd_mul"), l0, l1, l2);
+        mm->add_return({r});
+        auto prog = migraphx::parse_onnx("scatternd_mul_test.onnx");
+
+        EXPECT(p == prog);
+    }
+}
+
 TEST_CASE(selu_test)
 {
     migraphx::program p;

test/onnx/scatternd_add_test.onnx

+scatternd_add_test:Î

+@
+data
+indices
+updatesoutput"	ScatterND*
+	reduction"add 
scatternd_add_testZ
+data
+

+
+
+Z
+indices
+
+
+

+Z
+updates
+

+
+

+b
+output
+

+
+
+B
\ No newline at end of file

test/onnx/scatternd_mul_test.onnx

+scatternd_mul_test:

+@
+data
+indices
+updatesoutput"	ScatterND*
+	reduction"mul
scatternd_mul_testZ
+data
+

+
+
+Z
+indices
+
+
+

+Z
+updates
+

+
+

+b
+output
+

+
+
+B
\ No newline at end of file

test/onnx/scatternd_test.onnx

+scatternd_test:

++
+data
+indices
+updatesoutput"	ScatterNDscatternd_testZ
+data
+

+
+
+Z
+indices
+
+
+

+Z
+updates
+

+
+

+b
+output
+

+
+
+B
\ No newline at end of file

test/op_shape_test.cpp

@@ -1432,6 +1432,29 @@ TEST_CASE(test_scalar_nelemnts)
     throws_shape(migraphx::make_op("scalar", {{"scalar_bcst_dims", {2, 3, 4, 5}}}), input);
 }
 
+TEST_CASE(test_scatternd)
+{
+    {
+        // k > r
+        auto dtype = migraphx::shape::float_type;
+        auto itype = migraphx::shape::int64_type;
+        migraphx::shape ds{dtype, {8}};
+        migraphx::shape is{itype, {4, 2}};
+        migraphx::shape us{dtype, {4}};
+        throws_shape(migraphx::make_op("scatternd_none"), ds, is, us);
+    }
+
+    {
+        // update.lens != indices.lens[0:q-1] ++ data.lens[k:r-1]
+        auto dtype = migraphx::shape::float_type;
+        auto itype = migraphx::shape::int64_type;
+        migraphx::shape ds{dtype, {8}};
+        migraphx::shape is{itype, {4, 1}};
+        migraphx::shape us{dtype, {2, 2}};
+        throws_shape(migraphx::make_op("scatternd_none"), ds, is, us);
+    }
+}
+
 TEST_CASE(test_squeeze)
 {
     migraphx::shape s1{migraphx::shape::float_type, {4, 1, 3, 1, 3}};

test/py/onnx_backend_test.py

@@ -122,6 +122,7 @@ def create_backend_test(testname=None, target_device=None):
         backend_test.include(r'.*test_hardswish.*')
         backend_test.include(r'.*test_identity.*')
         backend_test.include(r'.*test_if.*')
+        backend_test.include(r'.*test_isnan.*')
         backend_test.include(r'.*test_LeakyReLU*')
         backend_test.include(r'.*test_leakyrelu.*')
         backend_test.include(r'.*test_less.*')
@@ -270,7 +271,6 @@ def create_backend_test(testname=None, target_device=None):
         backend_test.exclude(r'test_identity_sequence_cpu')
         backend_test.exclude(r'test_maxpool_2d_uint8_cpu')
         backend_test.exclude(r'test_negative_log_likelihood_loss_*')
-        backend_test.exclude(r'test_scatternd_*')
 
         # all reduce ops have dynamic axes inputs
         backend_test.exclude(r'test_size_cpu')

test/ref_ops_nonstd_shape_test.cpp

@@ -47,6 +47,26 @@ TEST_CASE(argmin_test_nonstd_shape)
     EXPECT(migraphx::verify_range(result_vec, res_gold_vec));
 }
 
+TEST_CASE(isnan_broadcast_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape s0{migraphx::shape::float_type, {3}};
+    migraphx::shape s1{migraphx::shape::float_type, {3, 2}};
+    auto nan_val             = std::numeric_limits<float>::quiet_NaN();
+    std::vector<float> data0 = {1.2, 5.2, nan_val};
+    auto l0                  = mm->add_literal(migraphx::literal{s0, data0});
+    auto l1                  = mm->add_instruction(
+        migraphx::make_op("broadcast", {{"axis", 0}, {"out_lens", s1.lens()}}), l0);
+    mm->add_instruction(migraphx::make_op("isnan"), l1);
+    p.compile(migraphx::ref::target{});
+    auto result = p.eval({}).back();
+    std::vector<float> results_vector;
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> correct = {0, 0, 0, 0, 1, 1};
+    EXPECT(migraphx::verify_range(results_vector, correct));
+}
+
 TEST_CASE(squeeze_transpose_test)
 {
     migraphx::program p;