batch_norm_op_xpu.cc 16.5 KB
Newer Older
1
/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
2 3 4 5 6 7 8 9 10 11 12 13 14 15
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#ifdef PADDLE_WITH_XPU

16 17
#include <iterator>
#include <vector>
18

19 20
#include "paddle/fluid/operators/batch_norm_op.h"

21 22 23 24 25 26 27 28 29
namespace paddle {
namespace operators {

using Tensor = framework::Tensor;
using DDim = framework::DDim;

template <typename DeviceContext, typename T>
class BatchNormXPUKernel : public framework::OpKernel<T> {
 public:
30
  void Compute(const framework::ExecutionContext &ctx) const override {
31
    const auto epsilon = ctx.Attr<float>("epsilon");
32
    float momentum = ctx.Attr<float>("momentum");
33 34 35 36
    const auto is_test = ctx.Attr<bool>("is_test");
    const auto use_global_stats = ctx.Attr<bool>("use_global_stats");
    const auto trainable_stats = ctx.Attr<bool>("trainable_statistics");
    bool test_mode = is_test && (!trainable_stats);
37

38
    bool global_stats = test_mode || use_global_stats;
39
    const auto &data_layout_str = ctx.Attr<std::string>("data_layout");
40
    const auto data_layout = framework::StringToDataLayout(data_layout_str);
41 42 43 44 45 46 47
    PADDLE_ENFORCE_EQ(data_layout_str == "NCHW" || data_layout_str == "NHWC",
                      true,
                      platform::errors::InvalidArgument(
                          "The 'data_layout' attribute must be NCHW or NHWC. "
                          "But recevived 'data_layout' is [%s].",
                          data_layout_str));

48 49
    const auto *x = ctx.Input<Tensor>("X");
    const auto &x_dims = x->dims();
50
    PADDLE_ENFORCE_EQ(
51 52
        x_dims.size() >= 2 && x_dims.size() <= 5,
        true,
53 54 55 56 57
        platform::errors::InvalidArgument(
            "The size of input's dimensions should be between 2 and 5"
            "But received: the size of input's dimensions is [%d]",
            x_dims.size()));

58
    int N = -1, C = -1, H = -1, W = -1, D = -1;
59
    ExtractNCWHD(x_dims, data_layout, &N, &C, &H, &W, &D);
60 61 62 63
    N = (N == 0) ? 1 : N;
    C = (C == 0) ? 1 : C;
    H = (H == 0) ? 1 : H;
    W = (W == 0) ? 1 : W;
64

65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84
    const auto *scale = ctx.Input<Tensor>("Scale");
    const auto *bias = ctx.Input<Tensor>("Bias");
    const auto *x_data = x->data<T>();
    const auto *scale_data = scale->data<float>();
    const auto *bias_data = bias->data<float>();

    auto *y = ctx.Output<Tensor>("Y");
    auto *mean_out = ctx.Output<Tensor>("MeanOut");
    auto *variance_out = ctx.Output<Tensor>("VarianceOut");
    auto *saved_mean = ctx.Output<Tensor>("SavedMean");
    auto *saved_variance = ctx.Output<Tensor>("SavedVariance");

    // alloc memory
    auto *y_data = y->mutable_data<T>(ctx.GetPlace());
    mean_out->mutable_data<float>(ctx.GetPlace());
    variance_out->mutable_data<float>(ctx.GetPlace());
    saved_mean->mutable_data<float>(ctx.GetPlace());
    saved_variance->mutable_data<float>(ctx.GetPlace());

    auto &dev_ctx = ctx.template device_context<DeviceContext>();
85
    bool is_nchw = data_layout_str == "NCHW";
86

87
    if (!global_stats) {
88 89 90 91 92 93 94 95 96 97
      auto *mean_out_data = mean_out->data<float>();
      auto *variance_out_data = variance_out->data<float>();
      auto *saved_mean_data = saved_mean->data<float>();
      auto *saved_variance_data = saved_variance->data<float>();

      // if MomentumTensor is set, use MomentumTensor value, momentum
      // is only used in this training branch
      if (ctx.HasInput("MomentumTensor")) {
        const auto *mom_tensor = ctx.Input<Tensor>("MomentumTensor");
        Tensor mom_cpu;
98 99
        paddle::framework::TensorCopySync(
            *mom_tensor, platform::CPUPlace(), &mom_cpu);
100 101
        momentum = mom_tensor->data<float>()[0];
      }
102

103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120
      int r = xpu::batch_norm<T>(dev_ctx.x_context(),
                                 x_data,
                                 y_data,
                                 N,
                                 C,
                                 H,
                                 W,
                                 epsilon,
                                 momentum,
                                 scale_data,
                                 bias_data,
                                 saved_mean_data,
                                 saved_variance_data,
                                 mean_out_data,
                                 variance_out_data,
                                 is_nchw);
      PADDLE_ENFORCE_EQ(r,
                        xpu::Error_t::SUCCESS,
121 122
                        platform::errors::External(
                            "The batch_norm XPU API return wrong value[%d %s]",
123 124
                            r,
                            XPUAPIErrorMsg[r]));
125
    } else {
126 127 128 129
      const auto *mean = ctx.Input<Tensor>("Mean");
      const auto *variance = ctx.Input<Tensor>("Variance");
      const auto *mean_data = mean->data<float>();
      const auto *variance_data = variance->data<float>();
130 131 132 133 134 135 136 137 138 139 140 141 142
      int r = xpu::batch_norm_infer(dev_ctx.x_context(),
                                    x_data,
                                    y_data,
                                    N,
                                    C,
                                    H,
                                    W,
                                    epsilon,
                                    scale_data,
                                    bias_data,
                                    mean_data,
                                    variance_data,
                                    is_nchw);
143
      PADDLE_ENFORCE_EQ(
144 145
          r,
          xpu::Error_t::SUCCESS,
146
          platform::errors::External(
147 148
              "The batch_norm_infer XPU API return wrong value[%d %s]",
              r,
149
              XPUAPIErrorMsg[r]));
150 151 152 153
    }
  }
};

154
template <typename T>
155 156 157 158 159 160 161 162 163
static int calculate_inv_BN_Y(xpu::Context *ctx,
                              T *x,
                              const T *scale,
                              const T *bias,
                              const T *mean,
                              const T *variance,
                              const int N,
                              const int C,
                              const int M,
164
                              const T *y) {
165 166
  PADDLE_ENFORCE_EQ(x,
                    y,
167 168
                    platform::errors::InvalidArgument(
                        "X and Y should be inplaced in inplace mode"));
169 170 171 172 173 174
  std::vector<int> tensor_shape_vec({N, C, M});
  std::vector<int> array_shape_vec({1, C, 1});
  // y - bias
  int r1 =
      xpu::broadcast_sub<T>(ctx, bias, y, x, array_shape_vec, tensor_shape_vec);
  // (y - bias) / scale
175 176
  int r2 = xpu::broadcast_div<T>(
      ctx, scale, x, x, array_shape_vec, tensor_shape_vec);
177
  // (y - bias) / scale / variance
178 179
  int r3 = xpu::broadcast_div<T>(
      ctx, variance, x, x, array_shape_vec, tensor_shape_vec);
180 181 182 183 184 185 186 187
  // (y - bias) / scale / variance + mean
  int r4 =
      xpu::broadcast_add<T>(ctx, mean, x, x, array_shape_vec, tensor_shape_vec);

  return r1 + r2 + r3 + r4;
}

template <typename T>
188 189 190 191 192 193
static int calculate_inv_var(xpu::Context *ctx,
                             const T *var,
                             const T epsilon,
                             const int C,
                             T *epsilon_data,
                             T *inv_var) {
194 195 196
  int r1 = constant(ctx, epsilon_data, 1, epsilon);
  std::vector<int> tensor_shape_vec({C});
  std::vector<int> array_shape_vec({1});
197 198
  int r2 = xpu::broadcast_add<T>(
      ctx, epsilon_data, var, inv_var, array_shape_vec, tensor_shape_vec);
199 200 201 202
  int r3 = xpu::rsqrt<T>(ctx, inv_var, inv_var, C);
  return r1 + r2 + r3;
}

203 204 205
template <typename DeviceContext, typename T>
class BatchNormGradXPUKernel : public framework::OpKernel<T> {
 public:
206 207 208 209 210 211 212 213 214
  void Compute(const framework::ExecutionContext &ctx) const override {
    const auto *d_y = ctx.Input<Tensor>(framework::GradVarName("Y"));
    const auto *scale = ctx.Input<Tensor>("Scale");
    const auto *bias = ctx.Input<Tensor>("Bias");

    const auto &data_layout_str = ctx.Attr<std::string>("data_layout");
    bool use_global_stats = ctx.Attr<bool>("use_global_stats");
    const bool is_test = ctx.Attr<bool>("is_test");
    const float epsilon = ctx.Attr<float>("epsilon");
215
    const auto data_layout = framework::StringToDataLayout(data_layout_str);
216

217 218 219 220 221 222 223
    PADDLE_ENFORCE_EQ(data_layout_str == "NCHW" || data_layout_str == "NHWC",
                      true,
                      platform::errors::InvalidArgument(
                          "The 'data_layout' attribute must be NCHW or NHWC. "
                          "But recevived 'data_layout' is [%s].",
                          data_layout_str));

224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240
    auto *d_x = ctx.Output<Tensor>(framework::GradVarName("X"));
    auto *d_scale = ctx.Output<Tensor>(framework::GradVarName("Scale"));
    auto *d_bias = ctx.Output<Tensor>(framework::GradVarName("Bias"));

    use_global_stats = is_test || use_global_stats;

    // batch_norm with inplace as false will take X as grad input, which
    // is same as cuDNN batch_norm backward calculation, batch_norm
    // with inplace as true only take Y as input and X should be calculate
    // by inverse operation of batch_norm on Y
    const Tensor *x;
    bool is_inplace;
    if (ctx.HasInput("Y")) {
      x = ctx.Input<Tensor>("Y");
      is_inplace = true;
      // if the input of batch norm is stop_gradient, d_x is null.
      if (d_x) {
241 242
        PADDLE_ENFORCE_EQ(d_x,
                          d_y,
243 244 245 246 247 248 249 250
                          platform::errors::InvalidArgument(
                              "X@GRAD and Y@GRAD not inplace in inplace mode"));
      }
    } else {
      x = ctx.Input<Tensor>("X");
      is_inplace = false;
      if (d_x) {
        PADDLE_ENFORCE_NE(
251 252
            d_x,
            d_y,
253 254
            platform::errors::InvalidArgument(
                "X@GRAD and Y@GRAD inplaced in non-inplace mode"));
255 256 257 258
      }
    }

    const auto &x_dims = x->dims();
259
    PADDLE_ENFORCE_EQ(
260 261
        x_dims.size() >= 2 && x_dims.size() <= 5,
        true,
262 263 264 265 266
        platform::errors::InvalidArgument(
            "The size of input's dimensions should be between 2 and 5"
            "But received: the size of input's dimensions is [%d]",
            x_dims.size()));

267
    int N = -1, C = -1, H = -1, W = -1, D = -1;
268
    ExtractNCWHD(x_dims, data_layout, &N, &C, &H, &W, &D);
269 270 271 272
    N = (N == 0) ? 1 : N;
    C = (C == 0) ? 1 : C;
    H = (H == 0) ? 1 : H;
    W = (W == 0) ? 1 : W;
273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290

    const auto *x_data = x->data<T>();
    const auto *d_y_data = d_y->data<T>();
    const auto *scale_data = scale->data<float>();

    // init output
    T *d_x_data = nullptr;
    T *d_bias_data = nullptr;
    T *d_scale_data = nullptr;
    if (d_x) {
      d_x_data = d_x->mutable_data<T>(ctx.GetPlace());
    }
    if (d_scale && d_bias) {
      d_scale_data = d_scale->mutable_data<float>(ctx.GetPlace());
      d_bias_data = d_bias->mutable_data<float>(ctx.GetPlace());
    }

    PADDLE_ENFORCE_EQ(
291 292
        scale->dims().size(),
        1UL,
293 294 295 296
        platform::errors::InvalidArgument(
            "The size of scale's dimensions must equal to 1. But received: "
            "the size of scale's dimensions is [%d], the dimensions of scale "
            "is [%s].",
297 298
            scale->dims().size(),
            scale->dims()));
299
    PADDLE_ENFORCE_EQ(
300 301
        scale->dims()[0],
        C,
302 303 304
        platform::errors::InvalidArgument(
            "The first dimension of scale must equal to Channels[%d]. But "
            "received: the first dimension of scale is [%d]",
305 306
            C,
            scale->dims()[0]));
307 308 309 310

    auto &dev_ctx = ctx.template device_context<DeviceContext>();
    xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());

311 312 313 314
    const auto *batch_mean = ctx.Input<Tensor>("SavedMean");
    const auto *batch_inv_std = ctx.Input<Tensor>("SavedVariance");
    const auto *global_mean = ctx.Input<Tensor>("Mean");
    const auto *global_var = ctx.Input<Tensor>("Variance");
315 316 317

    // TODO(guozibin): hadle the situation case of N * H * W = 1
    if (is_inplace) {
318 319 320 321 322
      float *global_inv_std_data = nullptr;
      if (use_global_stats) {
        global_inv_std_data =
            RAII_GUARD.alloc_l3_or_gm<float>(global_var->numel());
        float *epsilon_data = RAII_GUARD.alloc_l3_or_gm<float>(1);
323 324 325 326 327 328
        int r1 = calculate_inv_var(dev_ctx.x_context(),
                                   global_var->data<float>(),
                                   epsilon,
                                   C,
                                   epsilon_data,
                                   global_inv_std_data);
329
        PADDLE_ENFORCE_EQ(
330 331
            r1,
            XPU_SUCCESS,
332 333 334
            platform::errors::External("XPU API(batch_norm_grad "
                                       "calculate_inv_var function) "
                                       "return wrong value[%d %s]",
335 336
                                       r1,
                                       XPUAPIErrorMsg[r1]));
337
      }
338
      auto px = *x;
339 340 341 342
      auto *inv_std_data =
          use_global_stats ? global_inv_std_data : batch_inv_std->data<float>();
      auto mean_data = use_global_stats ? global_mean->data<float>()
                                        : batch_mean->data<float>();
343 344 345 346 347 348 349 350 351 352
      int r2 = calculate_inv_BN_Y(dev_ctx.x_context(),
                                  px.mutable_data<T>(ctx.GetPlace()),
                                  scale->data<float>(),
                                  bias->data<float>(),
                                  mean_data,
                                  inv_std_data,
                                  N,
                                  C,
                                  H * W,
                                  x->data<T>());
353
      PADDLE_ENFORCE_EQ(
354 355
          r2,
          XPU_SUCCESS,
356 357 358
          platform::errors::External("XPU API(batch_norm_grad "
                                     "calculate_inv_BN_Y function) "
                                     "return wrong value[%d %s]",
359 360
                                     r2,
                                     XPUAPIErrorMsg[r2]));
361
    }
362

363 364 365
    int r3;
    bool is_nchw = data_layout_str == "NCHW";
    if (use_global_stats) {
366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382
      r3 = xpu::batch_norm_grad<T>(dev_ctx.x_context(),
                                   x_data,
                                   d_y_data,
                                   d_x_data,
                                   N,
                                   C,
                                   H,
                                   W,
                                   scale_data,
                                   nullptr,
                                   nullptr,
                                   d_scale_data,
                                   d_bias_data,
                                   is_nchw,
                                   global_mean->data<float>(),
                                   global_var->data<float>(),
                                   epsilon);
383 384 385 386 387 388 389 390 391 392
    } else {
      if (!d_x) {
        d_x_data = RAII_GUARD.alloc_l3_or_gm<T>(x->numel());
      }
      if (!d_scale) {
        d_scale_data = RAII_GUARD.alloc_l3_or_gm<float>(C);
      }
      if (!d_bias_data) {
        d_bias_data = RAII_GUARD.alloc_l3_or_gm<float>(C);
      }
393 394 395 396 397 398 399 400 401 402 403 404 405 406
      r3 = xpu::batch_norm_grad<T>(dev_ctx.x_context(),
                                   x_data,
                                   d_y_data,
                                   d_x_data,
                                   N,
                                   C,
                                   H,
                                   W,
                                   scale_data,
                                   batch_mean->data<float>(),
                                   batch_inv_std->data<float>(),
                                   d_scale_data,
                                   d_bias_data,
                                   is_nchw);
407
    }
408
    PADDLE_ENFORCE_EQ(
409 410
        r3,
        XPU_SUCCESS,
411 412
        platform::errors::External("XPU API(batch_norm_grad) return "
                                   "wrong value[%d %s]",
413 414
                                   r3,
                                   XPUAPIErrorMsg[r3]));
415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;

REGISTER_OP_XPU_KERNEL(
    batch_norm,
    ops::BatchNormXPUKernel<paddle::platform::XPUDeviceContext, float>);
REGISTER_OP_XPU_KERNEL(
    batch_norm_grad,
    ops::BatchNormGradXPUKernel<paddle::platform::XPUDeviceContext, float>);

#endif  // PADDLE_WITH_XPU