add int32 support to argmin/argmax (#2288)

src/ngraph/runtime/cpu/builder/argmax.cpp

@@ -108,6 +108,34 @@ namespace ngraph
                         };
                     }
                 }
+                else if (element_type == element::i32)
+                {
+                    if (is_int64)
+                    {
+                        std::function<decltype(runtime::cpu::kernel::argmax<int, int64_t, 1>)>
+                            kernel;
+
+                        SELECT_RANK2(
+                            kernel, int, int64_t, in_shape.size(), runtime::cpu::kernel::argmax);
+
+                        functor = [&, kernel, in_shape, out_shape, axis](
+                            CPURuntimeContext* ctx, CPUExecutionContext* ectx) {
+                            kernel(arg_tensor, out_tensor, in_shape, out_shape, axis, ectx->arena);
+                        };
+                    }
+                    else
+                    {
+                        std::function<decltype(runtime::cpu::kernel::argmax<int, int, 1>)> kernel;
+
+                        SELECT_RANK2(
+                            kernel, int, int, in_shape.size(), runtime::cpu::kernel::argmax);
+
+                        functor = [&, kernel, in_shape, out_shape, axis](
+                            CPURuntimeContext* ctx, CPUExecutionContext* ectx) {
+                            kernel(arg_tensor, out_tensor, in_shape, out_shape, axis, ectx->arena);
+                        };
+                    }
+                }
                 else
                 {
                     throw ngraph_error("Unsupported type in CPU Builder for ArgMax");

src/ngraph/runtime/cpu/builder/argmin.cpp

@@ -108,6 +108,34 @@ namespace ngraph
                         };
                     }
                 }
+                else if (element_type == element::i32)
+                {
+                    if (is_int64)
+                    {
+                        std::function<decltype(runtime::cpu::kernel::argmin<int, int64_t, 1>)>
+                            kernel;
+
+                        SELECT_RANK2(
+                            kernel, int, int64_t, in_shape.size(), runtime::cpu::kernel::argmin);
+
+                        functor = [&, kernel, in_shape, out_shape, axis](
+                            CPURuntimeContext* ctx, CPUExecutionContext* ectx) {
+                            kernel(arg_tensor, out_tensor, in_shape, out_shape, axis, ectx->arena);
+                        };
+                    }
+                    else
+                    {
+                        std::function<decltype(runtime::cpu::kernel::argmin<int, int, 1>)> kernel;
+
+                        SELECT_RANK2(
+                            kernel, int, int, in_shape.size(), runtime::cpu::kernel::argmin);
+
+                        functor = [&, kernel, in_shape, out_shape, axis](
+                            CPURuntimeContext* ctx, CPUExecutionContext* ectx) {
+                            kernel(arg_tensor, out_tensor, in_shape, out_shape, axis, ectx->arena);
+                        };
+                    }
+                }
                 else
                 {
                     throw ngraph_error("Unsupported type in CPU Builder for ArgMin");