[phi] Add output defs for argsort kernel (#51407)

* Add output defs for argsort kernel

* Update argsort_kernel.cc

* Update argsort_kernel.cu

* Update argsort_kernel.cc

paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc

@@ -54,7 +54,6 @@ static std::set<std::string> OpsNeedSetOutputDtypeWhenRegisterPhiKernel = {
     "adam",
     "adamw",
     "any_raw",
-    "arg_sort",
     "clip_by_norm",
     "eig_grad",
     "eigh",

paddle/phi/kernels/cpu/argsort_kernel.cc

@@ -154,4 +154,5 @@ void ArgsortKernel(const Context& dev_ctx,
 
 PD_REGISTER_KERNEL(
     argsort, CPU, ALL_LAYOUT, phi::ArgsortKernel, float, double, int, int64_t) {
+  kernel->OutputAt(1).SetDataType(phi::DataType::INT64);
 }

paddle/phi/kernels/gpu/argsort_kernel.cu

@@ -519,4 +519,6 @@ PD_REGISTER_KERNEL(argsort,
                    double,
                    int,
                    int64_t,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16) {
+  kernel->OutputAt(1).SetDataType(phi::DataType::INT64);
+}

paddle/phi/kernels/xpu/argsort_kernel.cc

@@ -248,4 +248,6 @@ PD_REGISTER_KERNEL(argsort,
                    float,
                    int,
                    int64_t,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16) {
+  kernel->OutputAt(1).SetDataType(phi::DataType::INT64);
+}