[XPU] use int64_t in c_softmax (#62815)

PaddlePaddle · Mar 19, 2024 · 6307361 · 6307361
1 parent b67004f
commit 6307361
Showing 1 changed file with 14 additions and 14 deletions.
diff --git a/paddle/fluid/operators/collective/c_softmax_with_cross_entropy_op_xpu.cc b/paddle/fluid/operators/collective/c_softmax_with_cross_entropy_op_xpu.cc
@@ -83,8 +83,8 @@ struct CSoftmaxWithCrossEntropyProcessGroupFunctor<phi::XPUContext, T> {
     const auto& logits_dims = logits->dims();
 
     const int axis = logits_dims.size() - 1;
-    const int N = phi::funcs::SizeToAxis(axis, logits_dims);
-    const int D = phi::funcs::SizeFromAxis(axis, logits_dims);
+    const int64_t N = phi::funcs::SizeToAxis(axis, logits_dims);
+    const int64_t D = phi::funcs::SizeFromAxis(axis, logits_dims);
 
     phi::DenseTensor logits_2d, softmax_2d;
     framework::TensorCopy(
@@ -151,8 +151,8 @@ struct CSoftmaxWithCrossEntropyProcessGroupFunctor<phi::XPUContext, T> {
         N,
         0.0);
     PADDLE_ENFORCE_XDNN_SUCCESS(ret, "constant");
-    const int start_index = rank * D;
-    const int end_index = start_index + D;
+    const int64_t start_index = rank * D;
+    const int64_t end_index = start_index + D;
     const auto& label_type = framework::TransToProtoVarType(labels->dtype());
     if (label_type == framework::proto::VarType::INT32) {
       ret = xpu::mask_label_by_index<XPUType, int32_t>(
@@ -224,7 +224,7 @@ struct CSoftmaxWithCrossEntropyProcessGroupFunctor<phi::XPUContext, T> {
     opts.reduce_op = distributed::ReduceOp::SUM;
     pg->AllReduce(in_out, in_out, opts)->Synchronize();
 
-    int dims[4] = {N, D, N, 1};
+    int64_t dims[4] = {N, D, N, 1};
     ret = xpu::broadcast_div<XPUType>(
         dev_ctx.x_context(),
         reinterpret_cast<const XPUType*>(softmax_2d.data<T>()),
@@ -313,8 +313,8 @@ struct CSoftmaxWithCrossEntropyFunctor<phi::XPUContext, T> {
     const auto& logits_dims = logits->dims();
 
     const int axis = logits_dims.size() - 1;
-    const int N = phi::funcs::SizeToAxis(axis, logits_dims);
-    const int D = phi::funcs::SizeFromAxis(axis, logits_dims);
+    const int64_t N = phi::funcs::SizeToAxis(axis, logits_dims);
+    const int64_t D = phi::funcs::SizeFromAxis(axis, logits_dims);
 
     phi::DenseTensor logits_2d, softmax_2d;
     framework::TensorCopy(
@@ -390,8 +390,8 @@ struct CSoftmaxWithCrossEntropyFunctor<phi::XPUContext, T> {
         N,
         0.0);
     PADDLE_ENFORCE_XDNN_SUCCESS(ret, "constant");
-    const int start_index = rank * D;
-    const int end_index = start_index + D;
+    const int64_t start_index = rank * D;
+    const int64_t end_index = start_index + D;
     const auto& label_type = framework::TransToProtoVarType(labels->dtype());
     if (label_type == framework::proto::VarType::INT32) {
       ret = xpu::mask_label_by_index<XPUType, int32_t>(
@@ -485,7 +485,7 @@ struct CSoftmaxWithCrossEntropyFunctor<phi::XPUContext, T> {
     }
 
     {
-      int dims[4] = {N, D, N, 1};
+      int64_t dims[4] = {N, D, N, 1};
       ret = xpu::broadcast_div<XPUType>(
           dev_ctx.x_context(),
           reinterpret_cast<const XPUType*>(softmax_2d.data<T>()),
@@ -540,11 +540,11 @@ class CSoftmaxWithCrossEntropyGrad : public framework::OpKernel<T> {
     }
     const auto softmax_dims = softmax->dims();
     const int axis = softmax_dims.size() - 1;
-    const int N = phi::funcs::SizeToAxis(axis, softmax_dims);
-    const int D = phi::funcs::SizeFromAxis(axis, softmax_dims);
+    const int64_t N = phi::funcs::SizeToAxis(axis, softmax_dims);
+    const int64_t D = phi::funcs::SizeFromAxis(axis, softmax_dims);
 
-    const int start_index = rank * D;
-    const int end_index = start_index + D;
+    const int64_t start_index = rank * D;
+    const int64_t end_index = start_index + D;
     const auto& label_type = framework::TransToProtoVarType(labels->dtype());
 
     int ret = 0;