[XPU] support int64_t tril & triu (#68423)

runzhech · web-flow · commit a9b45a3d1805 · 2024-09-30T13:53:22.000+08:00
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -1233,33 +1233,45 @@ XPUOpMap& get_kl3_ops() {
       {"tril_triu",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
+                     phi::DataType::INT64,
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16,
                      phi::DataType::BOOL})},
       {"tril",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
+                     phi::DataType::INT64,
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16,
                      phi::DataType::BOOL})},
       {"triu",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
+                     phi::DataType::INT64,
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16,
                      phi::DataType::BOOL})},
       {"tril_triu_grad",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::INT64,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::BOOL})},
       {"tril_grad",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::INT64,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::BOOL})},
       {"triu_grad",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::INT64,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::BOOL})},
       {"tile",
        XPUKernelSet({phi::DataType::INT32,
                      phi::DataType::INT64,
diff --git a/paddle/phi/kernels/xpu/tril_triu_grad_kernel.cc b/paddle/phi/kernels/xpu/tril_triu_grad_kernel.cc
@@ -69,19 +69,28 @@ PD_REGISTER_KERNEL(tril_grad,
                    ALL_LAYOUT,
                    phi::TrilGradKernel,
                    int,
+                   int64_t,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   bool) {}
 PD_REGISTER_KERNEL(triu_grad,
                    XPU,
                    ALL_LAYOUT,
                    phi::TriuGradKernel,
                    int,
+                   int64_t,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   bool) {}
 PD_REGISTER_KERNEL(tril_triu_grad,
                    XPU,
                    ALL_LAYOUT,
                    phi::TrilTriuGradKernel,
                    int,
+                   int64_t,
                    float,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   bool) {}
diff --git a/paddle/phi/kernels/xpu/tril_triu_kernel.cc b/paddle/phi/kernels/xpu/tril_triu_kernel.cc
@@ -69,6 +69,7 @@ PD_REGISTER_KERNEL(tril_triu,
                    ALL_LAYOUT,
                    phi::TrilTriuKernel,
                    int,
+                   int64_t,
                    float,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
@@ -78,6 +79,7 @@ PD_REGISTER_KERNEL(tril,
                    ALL_LAYOUT,
                    phi::TrilKernel,
                    int,
+                   int64_t,
                    float,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
@@ -87,6 +89,7 @@ PD_REGISTER_KERNEL(triu,
                    ALL_LAYOUT,
                    phi::TriuKernel,
                    int,
+                   int64_t,
                    float,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
