coalesce_tensor_op.cc 21.9 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14
// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
//
// 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.

15
#include <sstream>
16 17
#include <vector>
#include "paddle/fluid/framework/op_registry.h"
18
#include "paddle/fluid/framework/op_version_registry.h"
19 20 21
#include "paddle/fluid/framework/operator.h"
#include "paddle/fluid/framework/var_type.h"
#include "paddle/fluid/operators/math/math_function.h"
22
#include "paddle/fluid/platform/device_memory_aligment.h"
23
#ifdef PADDLE_WITH_ASCEND_CL
24
#include "paddle/fluid/platform/device/npu/npu_op_runner.h"
25
#endif
26 27 28 29

namespace paddle {
namespace operators {

30 31 32
template <typename DeviceContext>
struct FillConstantVisitor {
  FillConstantVisitor(const DeviceContext &dev_ctx,
33 34 35 36 37 38 39 40
                      framework::LoDTensor *tensor, const float value,
                      framework::proto::VarType::Type dtype,
                      const framework::ExecutionContext &context)
      : dev_ctx_(dev_ctx),
        tensor_(tensor),
        value_(value),
        dtype_(dtype),
        context_(context) {}
41 42 43 44 45 46 47 48 49 50 51 52 53 54 55

  template <typename T>
  void apply(typename std::enable_if<std::is_same<T, int8_t>::value ||
                                     std::is_same<T, int16_t>::value>::type * =
                 nullptr) const {
    PADDLE_THROW(platform::errors::InvalidArgument(
        "Not support data type for set_constant attr"));
  }

  template <typename T>
  void apply(typename std::enable_if<!(std::is_same<T, int8_t>::value ||
                                       std::is_same<T, int16_t>::value)>::type
                 * = nullptr) const {
#ifdef PADDLE_WITH_ASCEND_CL
    if (platform::is_npu_place(dev_ctx_.GetPlace())) {
56 57 58 59 60 61 62 63 64 65 66
      Tensor tensor_tmp(dtype_);
      tensor_tmp.mutable_data<T>({1}, context_.GetPlace());
      FillNpuTensorWithConstant<T>(&tensor_tmp, static_cast<T>(value_));

      const auto &runner =
          NpuOpRunner("FillD", {tensor_tmp}, {*tensor_},
                      {{"dims", framework::vectorize(tensor_->dims())}});
      auto stream =
          context_.template device_context<paddle::platform::NPUDeviceContext>()
              .stream();
      runner.Run(stream);
67 68 69 70 71 72 73 74 75 76 77 78 79
    } else {
      math::SetConstant<DeviceContext, T> set_constant;
      set_constant(dev_ctx_, tensor_, static_cast<T>(value_));
    }
#else
    math::SetConstant<DeviceContext, T> set_constant;
    set_constant(dev_ctx_, tensor_, static_cast<T>(value_));
#endif
  }

  const DeviceContext &dev_ctx_;
  framework::LoDTensor *tensor_;
  float value_;
80 81
  framework::proto::VarType::Type dtype_;
  const framework::ExecutionContext &context_;
82 83
};

84
template <typename DeviceContext, typename T>
85
class CoalesceTensorOpKernel : public framework::OpKernel<T> {
86 87
 public:
  void Compute(const framework::ExecutionContext &context) const override {
H
hong 已提交
88 89
    auto in_var_names = context.InputNames("Input");
    auto out_var_names = context.OutputNames("Output");
90 91
    const auto &in_tensors = context.MultiInput<framework::LoDTensor>("Input");
    auto out_tensors = context.MultiOutput<framework::LoDTensor>("Output");
92

93
    PADDLE_ENFORCE_GT(in_var_names.size(), static_cast<size_t>(0),
94 95 96 97 98 99 100 101
                      platform::errors::InvalidArgument(
                          "The CoalesceTensor operator has no input."));
    PADDLE_ENFORCE_EQ(in_var_names.size(), out_var_names.size(),
                      platform::errors::InvalidArgument(
                          "The number of CoalesceTensor operator's input and "
                          "output is not match, "
                          "input number is %u, output number is %u.",
                          in_var_names.size(), out_var_names.size()));
102

103
    // Input & Output check: only support LoDTensor
104 105
    bool has_not_init_in_vars = false;
    for (size_t i = 0; i < in_tensors.size(); ++i) {
106
      PADDLE_ENFORCE_NOT_NULL(
107 108
          in_tensors[i], platform::errors::InvalidArgument(
                             "The %d-th input tensor cannot be nullptr.", i));
109
      PADDLE_ENFORCE_NOT_NULL(
110 111 112 113 114 115 116 117 118 119 120 121 122
          out_tensors[i], platform::errors::InvalidArgument(
                              "The %d-th output tensor cannot be nullptr.", i));
      if (!in_tensors[i]->IsInitialized()) {
        has_not_init_in_vars = true;
      }
    }

    if (has_not_init_in_vars) {
      const auto &concated_shapes =
          context.Attr<std::vector<int64_t>>("concated_shapes");
      const auto &concated_ranks =
          context.Attr<std::vector<int64_t>>("concated_ranks");
      PADDLE_ENFORCE_EQ(concated_ranks.size(), out_tensors.size(),
123
                        platform::errors::InvalidArgument(
124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153
                            "The attribute(concated_ranks) length must be "
                            "equal to the output tensor number."));
      int64_t accumulated_ranks = 0;
      for (size_t i = 0; i < in_tensors.size(); ++i) {
        framework::DDim dims(concated_shapes.data() + accumulated_ranks,
                             concated_ranks[i]);
        if (!in_tensors[i]->IsInitialized()) {
          PADDLE_ENFORCE_EQ(
              in_tensors[i], out_tensors[i],
              platform::errors::InvalidArgument(
                  "The %d-th output tensor and %d-th input tensor when the "
                  "%d-th input tensor is not initialized.",
                  i, i, i));
          out_tensors[i]->Resize(dims);
        } else {
          PADDLE_ENFORCE_EQ(
              in_tensors[i]->dims(), dims,
              platform::errors::InvalidArgument(
                  "The %d-th input tensor shape does not match the "
                  "attribute(concated_shapes) and "
                  "attribute(concated_ranks).",
                  i));
        }
        accumulated_ranks += concated_ranks[i];
        PADDLE_ENFORCE_LE(accumulated_ranks, concated_shapes.size(),
                          platform::errors::InvalidArgument(
                              "The attribute(concated_shapes) and "
                              "attribute(concated_ranks) do not match."));
      }
      PADDLE_ENFORCE_EQ(accumulated_ranks, concated_shapes.size(),
154
                        platform::errors::InvalidArgument(
155 156
                            "The attribute(concated_shapes) and "
                            "attribute(concated_ranks) do not match."));
157 158
    }

159
    bool use_align = context.Attr<bool>("use_align");
160
    auto align_size = context.Attr<int>("align_size");
161
    auto size_of_dtype = context.Attr<int>("user_defined_size_of_dtype");
162 163 164

    if (context.Attr<bool>("check_name")) {
      for (size_t i = 0; i < in_var_names.size(); ++i) {
165 166
        PADDLE_ENFORCE_EQ(
            in_var_names[i], out_var_names[i],
167 168 169 170
            platform::errors::InvalidArgument(
                "The input and output variable of CoalesceTensor operator is "
                "different, %dth input is %s, %dth output is %s.",
                i, in_var_names[i], i, out_var_names[i]));
171 172 173 174
      }
    } else {
      // Init the output as input
      for (size_t i = 0; i < in_tensors.size(); ++i) {
175
        out_tensors[i]->Resize(in_tensors[i]->dims());
176 177 178 179 180 181 182
      }
    }

    auto &dev_ctx = context.template device_context<DeviceContext>();

    // Get numel and dtype
    size_t numel = 0;
183 184
    auto dtype = static_cast<framework::proto::VarType::Type>(
        context.Attr<int>("dtype"));
185 186 187
    if (size_of_dtype == -1) {
      size_of_dtype = framework::SizeOfType(dtype);
    }
188
    GetMemSizeAndDtype(in_tensors, in_var_names, &numel, size_of_dtype,
189
                       context.GetPlace(), use_align, align_size);
190 191 192

    // Alloc the continuous space
    auto fused_tensor = context.Output<framework::LoDTensor>("FusedOutput");
193 194 195 196 197
    void *fused_tensor_ptr =
        fused_tensor
            ->Resize(framework::make_ddim({static_cast<int64_t>(numel)}))
            .mutable_data(context.GetPlace(), dtype);
    VLOG(10) << "Fused tensor addr " << fused_tensor_ptr;
198 199

    // Init the continuous space
C
chengduo 已提交
200
    size_t offset = 0;
201
    if (context.Attr<bool>("copy_data")) {
202 203 204 205 206 207
#ifdef PADDLE_WITH_ASCEND_CL
      framework::VisitDataType(
          dtype,
          FillConstantVisitor<DeviceContext>(
              dev_ctx, fused_tensor, static_cast<float>(0.0), dtype, context));
#endif
208
      for (size_t i = 0; i < in_var_names.size(); ++i) {
C
chengduo 已提交
209 210 211 212
        size_t len = static_cast<size_t>(in_tensors[i]->numel());
        auto sub_tensor = fused_tensor->Slice(
            static_cast<int64_t>(offset), static_cast<int64_t>(offset + len));
        framework::TensorCopy(*in_tensors[i], context.GetPlace(), dev_ctx,
213
                              &sub_tensor);
C
chengduo 已提交
214

215 216 217 218 219
        offset += use_align
                      ? platform::Alignment(len * size_of_dtype,
                                            context.GetPlace(), align_size) /
                            size_of_dtype
                      : len;
220 221
      }
    } else if (context.Attr<bool>("set_constant")) {
222 223
      framework::VisitDataType(
          dtype, FillConstantVisitor<DeviceContext>(
224 225
                     dev_ctx, fused_tensor, context.Attr<float>("constant"),
                     dtype, context));
226 227 228 229 230 231 232 233 234 235
    } else if (context.Attr<bool>("persist_output")) {
      for (size_t i = 0; i < out_var_names.size(); ++i) {
        size_t len = static_cast<size_t>(out_tensors[i]->numel());
        auto sub_tensor = fused_tensor->Slice(
            static_cast<int64_t>(offset), static_cast<int64_t>(offset + len));
        // some var may not persistable, or persistable var may not init
        if (out_tensors[i]->IsInitialized()) {
          framework::TensorCopy(*out_tensors[i], context.GetPlace(), dev_ctx,
                                &sub_tensor);
        }
236 237 238 239 240
        offset += use_align
                      ? platform::Alignment(len * size_of_dtype,
                                            context.GetPlace(), align_size) /
                            size_of_dtype
                      : len;
241
      }
242 243 244 245
    }

    // Make the outputs point to the continuous space.
    offset = 0;
246 247
    std::stringstream ss;
    ss << "alloc_space_for_vars: ";
248

249
    for (size_t i = 0; i < out_tensors.size(); ++i) {
C
chengduo 已提交
250
      size_t len = static_cast<size_t>(out_tensors[i]->numel());
251
      auto dim = out_tensors[i]->dims();
252
      VLOG(4) << len << " " << dim << " " << offset;
253
      out_tensors[i]
C
chengduo 已提交
254 255
          ->ShareDataWith(fused_tensor->Slice(
              static_cast<int64_t>(offset), static_cast<int64_t>(offset + len)))
256
          .Resize(dim);
257
      len = use_align
258 259
                ? platform::Alignment(len * size_of_dtype, context.GetPlace(),
                                      align_size) /
260 261
                      size_of_dtype
                : len;
262
      ss << "output(" << out_var_names[i] << ")  dim:(" << dim << ")"
263
         << " address: " << out_tensors[i]->data() << " len: " << len << ", ";
264
      offset += len;
265
    }
266 267 268 269 270 271
    PADDLE_ENFORCE_EQ(
        (int64_t)offset, fused_tensor->numel(),
        platform::errors::InvalidArgument(
            "The alloc_space_for_vars's offset: %s is unequal with "
            "fused_tensor's numel: %s.",
            offset, fused_tensor->numel()));
272
    VLOG(10) << ss.str();
273 274
  }

C
chengduo 已提交
275
 private:
276 277 278
  void GetMemSizeAndDtype(
      const std::vector<const framework::LoDTensor *> &lod_tensors,
      const std::vector<std::string> var_names, size_t *numel,
279
      const size_t &size_of_dtype, const platform::Place &place,
280
      const bool use_align = true, const int align_size = -1) const {
281 282 283 284 285 286
    PADDLE_ENFORCE_EQ(
        lod_tensors.size(), var_names.size(),
        platform::errors::InvalidArgument(
            "The number of input tensor and variable does not match, the "
            "number of input tensor is %u, the number of input variable is %u.",
            lod_tensors.size(), var_names.size()));
287
    *numel = 0;
288 289
    std::stringstream ss;
    ss << "alloc_space_for_vars: ";
290 291
    for (size_t i = 0; i < var_names.size(); ++i) {
      auto size = lod_tensors[i]->numel();
292 293 294 295
      PADDLE_ENFORCE_GT(
          size, 0,
          platform::errors::InvalidArgument(
              "The number of tensor `%s`'s elements is 0.", var_names[i]));
296 297 298 299 300 301
      auto len =
          use_align
              ? platform::Alignment(static_cast<size_t>(size) * size_of_dtype,
                                    place, align_size) /
                    size_of_dtype
              : static_cast<size_t>(size);
302 303
      const void *ptr =
          lod_tensors[i]->IsInitialized() ? lod_tensors[i]->data() : nullptr;
304
      VLOG(4) << size << " " << len;
305
      ss << "input(" << var_names[i] << ") dim:(" << lod_tensors[i]->dims()
306
         << ") "
307
         << " addres:" << ptr << " len: " << len << ", ";
308
      *numel += len;
309
    }
310
    VLOG(10) << ss.str();
311 312 313
  }
};

314
class CoalesceTensorOp : public framework::OperatorWithKernel {
315 316 317
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

318 319 320 321 322 323
  void InferShape(framework::InferShapeContext *ctx) const override {
    if (ctx->IsRuntime()) {
      return;
    }
    auto use_align = ctx->Attrs().Get<bool>("use_align");
    auto align_size = ctx->Attrs().Get<int>("align_size");
324
    auto size_of_dtype = ctx->Attrs().Get<int>("user_defined_size_of_dtype");
325 326 327

    auto dtype = static_cast<framework::proto::VarType::Type>(
        ctx->Attrs().Get<int>("dtype"));
328 329 330
    if (size_of_dtype == -1) {
      size_of_dtype = framework::SizeOfType(dtype);
    }
331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356

    auto alignment = [](size_t size, size_t align_size) {
      size_t remaining = size % align_size;
      auto aligned_size =
          remaining == 0 ? size : size + (align_size - remaining);
      VLOG(4) << remaining << " " << size << " " << align_size << " "
              << aligned_size;
      return aligned_size;
    };
    VLOG(4) << "align_size: " << align_size;
    if (use_align && align_size > 0) {
      int64_t numel = 0;
      auto dims = ctx->GetInputsDim("Input");
      for (const auto &dim : dims) {
        auto size = framework::product(dim);
        auto len = use_align
                       ? alignment(static_cast<size_t>(size) * size_of_dtype,
                                   align_size) /
                             size_of_dtype
                       : static_cast<size_t>(size);
        numel += len;
      }
      ctx->SetOutputDim("FusedOutput", framework::make_ddim({numel}));
      VLOG(4) << "FusedOutput size:" << framework::make_ddim({numel});
    }
  }
357 358

 protected:
359 360 361 362 363 364 365
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &context) const override {
    auto dtype = static_cast<framework::proto::VarType::Type>(
        context.Attr<int>("dtype"));
    return framework::OpKernelType(dtype, context.GetPlace());
  }

366 367 368 369 370 371 372
  framework::OpKernelType GetKernelTypeForVar(
      const std::string &var_name, const framework::Tensor &tensor,
      const framework::OpKernelType &expected_kernel_type) const override {
    return framework::OpKernelType(expected_kernel_type.data_type_,
                                   expected_kernel_type.place_,
                                   tensor.layout());
  }
373 374
};

375
class CoalesceTensorOpMaker : public framework::OpProtoAndCheckerMaker {
376 377 378 379
 public:
  void Make() override {
    AddInput("Input",
             "(vector<LoDTensor>) The input tensors of"
380
             " coalesce_tensor operator.")
381 382 383
        .AsDuplicable();
    AddOutput("Output",
              "(vector<LoDTensor>) The output "
384
              "tensors of coalesce_tensor operator. And the address "
385 386 387 388 389
              "of output tensors are continuous, they are sliced from the "
              "tensor of FusedOutput.")
        .AsDuplicable();
    AddOutput("FusedOutput",
              "(LoDTensor) The output tensor "
390
              "of coalesce_tensor operator. And the tensors of"
391
              " Output is sliced from the tensor of FusedOutput.");
392
    AddAttr<int>("dtype", "The output data type.");
393 394 395 396 397
    AddAttr<bool>("copy_data", "Whether to copy the Input value to Output.")
        .SetDefault(false);
    AddAttr<bool>("set_constant",
                  "Whether to set the Output with a constant value.")
        .SetDefault(false);
398 399 400
    AddAttr<bool>("persist_output",
                  "Whether to persist the original Output value.")
        .SetDefault(false);
401 402 403 404 405 406 407 408
    AddAttr<float>("constant",
                   "If set_constant is true, the constant value will be used "
                   "to set the Output.")
        .SetDefault(0.0);
    AddAttr<bool>("check_name",
                  "Whether to check the name of Input and Output to ensure "
                  "they are the same separately.")
        .SetDefault(false);
409 410 411 412
    AddAttr<bool>("use_align",
                  "Whether to consider memory chunk and take alignment into "
                  "account for inputs and outputs.")
        .SetDefault(true);
413 414
    AddAttr<int>("align_size", "The alignment size when use_align is True")
        .SetDefault(-1);
415 416 417 418 419 420 421 422 423
    AddAttr<int>("user_defined_size_of_dtype",
                 "The user defined size of dtype. This is used to coalesce "
                 "grad vars and merged_grad vars at the same time. For some "
                 "strategy, the dtype of fused_grad_vars and the dtype of "
                 "fused_grad_merged_vars are not identical, which will cause "
                 "the shape of these two coalesced vars are different. To "
                 "make sure the shape of these two vars are identical with "
                 "each other, this attr is added.")
        .SetDefault(-1);
424 425 426 427 428 429 430 431 432 433 434 435 436 437
    AddAttr<std::vector<int64_t>>(
        "concated_shapes",
        "The concated shapes of each shape of the input tensors. "
        "If any of the input tensors are not inited, this is used to "
        "init the output tensor shape, together with "
        "attribute(concated_ranks).")
        .SetDefault({});
    AddAttr<std::vector<int64_t>>(
        "concated_ranks",
        "The concated ranks of each rank of the input tensors. "
        "If any of the input tensors are not inited, this is used to "
        "init the output tensor shape, together with "
        "attribute(concated_shapes).")
        .SetDefault({});
438
    AddComment(R"DOC(
439
CoalesceTensor Operator.
440

441
coalesce_tensor is used to make the address of Output
442 443 444 445 446 447 448 449
continuous according to the Input. This Op will alloc a big tensor
according to the tensors of Input, the dtype is the same with those input tensors,
the size is the sum of those input tensors' numel, and the dim of the big
tensor is {sum(numel)}. And the big tensor is stored in FusedOutput.
The tensors of Output are sliced from the tensor of FusedOutput.
Note that, the dtype of Input should be the same, and the dim of Input
and Output should equal.
The tensors of Input and Output could be the same or different. And
450
coalesce_tensor allows copying the value of Input to Output, or
451 452
setting the Output with a constant value, or persist the original Output
value.
453 454 455 456 457 458 459 460

)DOC");
  }
};

}  // namespace operators
}  // namespace paddle

461 462
REGISTER_OPERATOR(coalesce_tensor, paddle::operators::CoalesceTensorOp,
                  paddle::operators::CoalesceTensorOpMaker);
463
namespace ops = paddle::operators;
464
namespace plat = paddle::platform;
465
REGISTER_OP_CPU_KERNEL(
466
    coalesce_tensor,
467 468 469
    ops::CoalesceTensorOpKernel<paddle::platform::CPUDeviceContext, int>,
    ops::CoalesceTensorOpKernel<paddle::platform::CPUDeviceContext, float>,
    ops::CoalesceTensorOpKernel<paddle::platform::CPUDeviceContext, double>);
470

471
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
472
REGISTER_OP_CUDA_KERNEL(
473
    coalesce_tensor,
474 475 476 477 478
    ops::CoalesceTensorOpKernel<paddle::platform::CUDADeviceContext,
                                plat::float16>,
    ops::CoalesceTensorOpKernel<paddle::platform::CUDADeviceContext, int>,
    ops::CoalesceTensorOpKernel<paddle::platform::CUDADeviceContext, float>,
    ops::CoalesceTensorOpKernel<paddle::platform::CUDADeviceContext, double>);
479
#endif
480

481 482 483 484 485 486 487 488 489 490
#if defined(PADDLE_WITH_ASCEND_CL)
REGISTER_OP_CUDA_KERNEL(
    coalesce_tensor,
    ops::CoalesceTensorOpKernel<paddle::platform::NPUDeviceContext,
                                plat::float16>,
    ops::CoalesceTensorOpKernel<paddle::platform::NPUDeviceContext, int>,
    ops::CoalesceTensorOpKernel<paddle::platform::NPUDeviceContext, float>,
    ops::CoalesceTensorOpKernel<paddle::platform::NPUDeviceContext, double>);
#endif

W
WangXi 已提交
491 492 493 494 495 496 497 498 499 500
#ifdef PADDLE_WITH_XPU
REGISTER_OP_XPU_KERNEL(
    coalesce_tensor,
    ops::CoalesceTensorOpKernel<paddle::platform::XPUDeviceContext,
                                plat::float16>,
    ops::CoalesceTensorOpKernel<paddle::platform::XPUDeviceContext, int>,
    ops::CoalesceTensorOpKernel<paddle::platform::XPUDeviceContext, float>,
    ops::CoalesceTensorOpKernel<paddle::platform::XPUDeviceContext, double>);
#endif

501 502 503 504 505 506 507 508 509 510
#if defined(PADDLE_WITH_ASCEND_CL)
REGISTER_OP_NPU_KERNEL(
    coalesce_tensor,
    ops::CoalesceTensorOpKernel<paddle::platform::CPUDeviceContext, int>,
    ops::CoalesceTensorOpKernel<paddle::platform::CPUDeviceContext, float>,
    ops::CoalesceTensorOpKernel<paddle::platform::CPUDeviceContext,
                                plat::float16>,
    ops::CoalesceTensorOpKernel<paddle::platform::CPUDeviceContext, double>);
#endif

511 512 513 514 515 516 517 518 519
REGISTER_OP_VERSION(coalesce_tensor)
    .AddCheckpoint(
        R"ROC(
              Upgrade coalesce_tensor: add a new attribute [use_align].)ROC",
        paddle::framework::compatible::OpVersionDesc().NewAttr(
            "use_align",
            "In order to optionally take memory alignment into account when "
            "coalescing tensors. The default value is true to be compatible "
            "with before.",
520 521 522 523 524 525 526 527 528 529 530
            true))
    .AddCheckpoint(
        R"ROC(
                Upgrade coalesce_tensor: add a new attribute [align_size].)ROC",
        paddle::framework::compatible::OpVersionDesc().NewAttr(
            "align_size",
            "In order to optionally take memory alignment into account when "
            "coalescing tensors. The default value is -1 and use the default "
            "align_size "
            "of each place to be compatible with before.",
            -1));