[NPU] support arg_max for int64 (#1374)

Author: wangna11BD

backends/npu/kernels/arg_min_max_kernel.cc

@@ -203,6 +203,7 @@ PD_REGISTER_PLUGIN_KERNEL(argmax,
                           ALL_LAYOUT,
                           custom_kernel::ArgMaxKernel,
                           int,
+                          int64_t,
                           float,
                           double,
                           phi::dtype::float16,

backends/npu/tests/unittests/test_arg_max_op_npu.py

@@ -382,6 +382,87 @@ def initTestCase(self):
         self.axis = 0
 
 
+# test argmax, dtype: int64
+class TestArgMaxInt64Case1(BaseTestCase):
+    def initTestCase(self):
+        self.op_type = "arg_max"
+        self.dims = (3, 4, 5)
+        self.dtype = "int64"
+        self.axis = -1
+
+
+class TestArgMaxInt64Case2(BaseTestCase):
+    def initTestCase(self):
+        self.op_type = "arg_max"
+        self.dims = (3, 4, 5)
+        self.dtype = "int64"
+        self.axis = 0
+
+
+class TestArgMaxInt64Case3(BaseTestCase):
+    def initTestCase(self):
+        self.op_type = "arg_max"
+        self.dims = (3, 4, 5)
+        self.dtype = "int64"
+        self.axis = 1
+
+
+class TestArgMaxInt64Case4(BaseTestCase):
+    def initTestCase(self):
+        self.op_type = "arg_max"
+        self.dims = (3, 4, 5)
+        self.dtype = "int64"
+        self.axis = 2
+
+
+class TestArgMaxInt64Case5(BaseTestCase):
+    def initTestCase(self):
+        self.op_type = "arg_max"
+        self.dims = (3, 4)
+        self.dtype = "int64"
+        self.axis = -1
+
+
+class TestArgMaxInt64Case6(BaseTestCase):
+    def initTestCase(self):
+        self.op_type = "arg_max"
+        self.dims = (3, 4)
+        self.dtype = "int64"
+        self.axis = 0
+
+
+class TestArgMaxInt64Case7(BaseTestCase):
+    def initTestCase(self):
+        self.op_type = "arg_max"
+        self.dims = (3, 4)
+        self.dtype = "int64"
+        self.axis = 1
+
+
+class TestArgMaxInt64Case8(BaseTestCase):
+    def initTestCase(self):
+        self.op_type = "arg_max"
+        self.dims = (1,)
+        self.dtype = "int64"
+        self.axis = 0
+
+
+class TestArgMaxInt64Case9(BaseTestCase):
+    def initTestCase(self):
+        self.op_type = "arg_max"
+        self.dims = (2,)
+        self.dtype = "int64"
+        self.axis = 0
+
+
+class TestArgMaxInt64Case10(BaseTestCase):
+    def initTestCase(self):
+        self.op_type = "arg_max"
+        self.dims = (3,)
+        self.dtype = "int64"
+        self.axis = 0
+
+
 class BaseTestComplex1_1(OpTest):
     def set_npu(self):
         self.__class__.use_custom_device = True