coalesce_tensor_op.cc 22.7 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
#include "paddle/fluid/framework/operator.h"
#include "paddle/fluid/framework/var_type.h"
21
#include "paddle/fluid/platform/device_memory_aligment.h"
22
#include "paddle/phi/kernels/funcs/math_function.h"
23
#ifdef PADDLE_WITH_ASCEND_CL
24
#include "paddle/fluid/platform/device/npu/npu_op_runner.h"
25
#endif
26
#include "paddle/fluid/framework/convert_utils.h"
27 28 29
#ifdef PADDLE_WITH_MLU
#include "paddle/fluid/operators/mlu/mlu_baseop.h"
#endif
30 31 32 33

namespace paddle {
namespace operators {

34 35 36
template <typename DeviceContext>
struct FillConstantVisitor {
  FillConstantVisitor(const DeviceContext &dev_ctx,
37 38 39 40 41 42 43 44
                      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) {}
45 46 47 48 49 50 51 52 53 54 55 56 57 58 59

  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())) {
60
      Tensor tensor_tmp(framework::TransToPhiDataType(dtype_));
61 62 63 64 65
      tensor_tmp.mutable_data<T>({1}, context_.GetPlace());
      FillNpuTensorWithConstant<T>(&tensor_tmp, static_cast<T>(value_));

      const auto &runner =
          NpuOpRunner("FillD", {tensor_tmp}, {*tensor_},
66
                      {{"dims", phi::vectorize(tensor_->dims())}});
67 68 69 70
      auto stream =
          context_.template device_context<paddle::platform::NPUDeviceContext>()
              .stream();
      runner.Run(stream);
71
    } else {
72
      phi::funcs::SetConstant<DeviceContext, T> set_constant;
73 74
      set_constant(dev_ctx_, tensor_, static_cast<T>(value_));
    }
75 76 77 78 79 80 81
#elif defined(PADDLE_WITH_MLU)
    if (platform::is_mlu_place(context_.GetPlace())) {
      FillMLUTensorWithHostValue<T>(context_, static_cast<T>(value_), tensor_);
    } else {
      phi::funcs::SetConstant<DeviceContext, T> set_constant;
      set_constant(dev_ctx_, tensor_, static_cast<T>(value_));
    }
82
#else
83
    phi::funcs::SetConstant<DeviceContext, T> set_constant;
84 85 86 87 88 89 90
    set_constant(dev_ctx_, tensor_, static_cast<T>(value_));
#endif
  }

  const DeviceContext &dev_ctx_;
  framework::LoDTensor *tensor_;
  float value_;
91 92
  framework::proto::VarType::Type dtype_;
  const framework::ExecutionContext &context_;
93 94
};

95
template <typename DeviceContext, typename T>
96
class CoalesceTensorOpKernel : public framework::OpKernel<T> {
97 98
 public:
  void Compute(const framework::ExecutionContext &context) const override {
H
hong 已提交
99 100
    auto in_var_names = context.InputNames("Input");
    auto out_var_names = context.OutputNames("Output");
101 102
    const auto &in_tensors = context.MultiInput<framework::LoDTensor>("Input");
    auto out_tensors = context.MultiOutput<framework::LoDTensor>("Output");
103

104
    PADDLE_ENFORCE_GT(in_var_names.size(), static_cast<size_t>(0),
105 106 107 108 109 110 111 112
                      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()));
113

114
    // Input & Output check: only support LoDTensor
115 116
    bool has_not_init_in_vars = false;
    for (size_t i = 0; i < in_tensors.size(); ++i) {
117
      PADDLE_ENFORCE_NOT_NULL(
118 119
          in_tensors[i], platform::errors::InvalidArgument(
                             "The %d-th input tensor cannot be nullptr.", i));
120
      PADDLE_ENFORCE_NOT_NULL(
121 122 123 124 125 126 127 128 129 130 131 132 133
          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(),
134
                        platform::errors::InvalidArgument(
135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164
                            "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(),
165
                        platform::errors::InvalidArgument(
166 167
                            "The attribute(concated_shapes) and "
                            "attribute(concated_ranks) do not match."));
168 169
    }

170
    bool use_align = context.Attr<bool>("use_align");
171
    auto align_size = context.Attr<int>("align_size");
172
    auto size_of_dtype = context.Attr<int>("user_defined_size_of_dtype");
173 174 175

    if (context.Attr<bool>("check_name")) {
      for (size_t i = 0; i < in_var_names.size(); ++i) {
176 177
        PADDLE_ENFORCE_EQ(
            in_var_names[i], out_var_names[i],
178 179 180 181
            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]));
182 183 184 185
      }
    } else {
      // Init the output as input
      for (size_t i = 0; i < in_tensors.size(); ++i) {
186
        out_tensors[i]->Resize(in_tensors[i]->dims());
187 188 189 190 191 192 193
      }
    }

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

    // Get numel and dtype
    size_t numel = 0;
194 195
    auto dtype = static_cast<framework::proto::VarType::Type>(
        context.Attr<int>("dtype"));
196 197 198
    if (size_of_dtype == -1) {
      size_of_dtype = framework::SizeOfType(dtype);
    }
199
    GetMemSizeAndDtype(in_tensors, in_var_names, &numel, size_of_dtype,
200
                       context.GetPlace(), use_align, align_size);
201 202 203

    // Alloc the continuous space
    auto fused_tensor = context.Output<framework::LoDTensor>("FusedOutput");
204
    void *fused_tensor_ptr =
205
        fused_tensor->Resize(phi::make_ddim({static_cast<int64_t>(numel)}))
206
            .mutable_data(context.GetPlace(),
207
                          framework::TransToPhiDataType(dtype));
208
    VLOG(10) << "Fused tensor addr " << fused_tensor_ptr;
209 210

    // Init the continuous space
C
chengduo 已提交
211
    size_t offset = 0;
212
    if (context.Attr<bool>("copy_data")) {
213 214 215 216 217 218
#ifdef PADDLE_WITH_ASCEND_CL
      framework::VisitDataType(
          dtype,
          FillConstantVisitor<DeviceContext>(
              dev_ctx, fused_tensor, static_cast<float>(0.0), dtype, context));
#endif
219
      for (size_t i = 0; i < in_var_names.size(); ++i) {
C
chengduo 已提交
220 221 222 223
        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,
224
                              &sub_tensor);
C
chengduo 已提交
225

226 227 228 229 230
        offset += use_align
                      ? platform::Alignment(len * size_of_dtype,
                                            context.GetPlace(), align_size) /
                            size_of_dtype
                      : len;
231 232
      }
    } else if (context.Attr<bool>("set_constant")) {
233 234
      framework::VisitDataType(
          dtype, FillConstantVisitor<DeviceContext>(
235 236
                     dev_ctx, fused_tensor, context.Attr<float>("constant"),
                     dtype, context));
237 238 239 240 241 242 243 244 245 246
    } 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);
        }
247 248 249 250 251
        offset += use_align
                      ? platform::Alignment(len * size_of_dtype,
                                            context.GetPlace(), align_size) /
                            size_of_dtype
                      : len;
252
      }
253 254 255 256
    }

    // Make the outputs point to the continuous space.
    offset = 0;
257 258
    std::stringstream ss;
    ss << "alloc_space_for_vars: ";
259

260
    for (size_t i = 0; i < out_tensors.size(); ++i) {
C
chengduo 已提交
261
      size_t len = static_cast<size_t>(out_tensors[i]->numel());
262
      auto dim = out_tensors[i]->dims();
263
      VLOG(4) << len << " " << dim << " " << offset;
264
      out_tensors[i]
C
chengduo 已提交
265 266
          ->ShareDataWith(fused_tensor->Slice(
              static_cast<int64_t>(offset), static_cast<int64_t>(offset + len)))
267
          .Resize(dim);
268
      len = use_align
269 270
                ? platform::Alignment(len * size_of_dtype, context.GetPlace(),
                                      align_size) /
271 272
                      size_of_dtype
                : len;
273
      ss << "output(" << out_var_names[i] << ")  dim:(" << dim << ")"
274
         << " address: " << out_tensors[i]->data() << " len: " << len << ", ";
275
      offset += len;
276
    }
277 278 279 280 281 282
    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()));
283
    VLOG(10) << ss.str();
284 285
  }

C
chengduo 已提交
286
 private:
287 288 289
  void GetMemSizeAndDtype(
      const std::vector<const framework::LoDTensor *> &lod_tensors,
      const std::vector<std::string> var_names, size_t *numel,
290
      const size_t &size_of_dtype, const platform::Place &place,
291
      const bool use_align = true, const int align_size = -1) const {
292 293 294 295 296 297
    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()));
298
    *numel = 0;
299 300
    std::stringstream ss;
    ss << "alloc_space_for_vars: ";
301 302
    for (size_t i = 0; i < var_names.size(); ++i) {
      auto size = lod_tensors[i]->numel();
303 304 305 306
      PADDLE_ENFORCE_GT(
          size, 0,
          platform::errors::InvalidArgument(
              "The number of tensor `%s`'s elements is 0.", var_names[i]));
307 308 309 310 311 312
      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);
313 314
      const void *ptr =
          lod_tensors[i]->IsInitialized() ? lod_tensors[i]->data() : nullptr;
315
      VLOG(4) << size << " " << len;
316
      ss << "input(" << var_names[i] << ") dim:(" << lod_tensors[i]->dims()
317
         << ") "
318
         << " addres:" << ptr << " len: " << len << ", ";
319
      *numel += len;
320
    }
321
    VLOG(10) << ss.str();
322 323 324
  }
};

325
class CoalesceTensorOp : public framework::OperatorWithKernel {
326 327 328
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

329 330 331 332 333 334
  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");
335
    auto size_of_dtype = ctx->Attrs().Get<int>("user_defined_size_of_dtype");
336 337 338

    auto dtype = static_cast<framework::proto::VarType::Type>(
        ctx->Attrs().Get<int>("dtype"));
339 340 341
    if (size_of_dtype == -1) {
      size_of_dtype = framework::SizeOfType(dtype);
    }
342 343 344 345 346 347 348 349 350 351 352 353 354 355

    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) {
356
        auto size = phi::product(dim);
357 358 359 360 361 362 363
        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;
      }
364 365
      ctx->SetOutputDim("FusedOutput", phi::make_ddim({numel}));
      VLOG(4) << "FusedOutput size:" << phi::make_ddim({numel});
366 367
    }
  }
368 369

 protected:
370 371 372 373 374 375 376
  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());
  }

377 378 379 380 381 382 383
  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());
  }
384 385
};

386
class CoalesceTensorOpMaker : public framework::OpProtoAndCheckerMaker {
387 388 389 390
 public:
  void Make() override {
    AddInput("Input",
             "(vector<LoDTensor>) The input tensors of"
391
             " coalesce_tensor operator.")
392 393 394
        .AsDuplicable();
    AddOutput("Output",
              "(vector<LoDTensor>) The output "
395
              "tensors of coalesce_tensor operator. And the address "
396 397 398 399 400
              "of output tensors are continuous, they are sliced from the "
              "tensor of FusedOutput.")
        .AsDuplicable();
    AddOutput("FusedOutput",
              "(LoDTensor) The output tensor "
401
              "of coalesce_tensor operator. And the tensors of"
402
              " Output is sliced from the tensor of FusedOutput.");
403
    AddAttr<int>("dtype", "The output data type.");
404 405 406 407 408
    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);
409 410 411
    AddAttr<bool>("persist_output",
                  "Whether to persist the original Output value.")
        .SetDefault(false);
412 413 414 415 416 417 418 419
    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);
420 421 422 423
    AddAttr<bool>("use_align",
                  "Whether to consider memory chunk and take alignment into "
                  "account for inputs and outputs.")
        .SetDefault(true);
424 425
    AddAttr<int>("align_size", "The alignment size when use_align is True")
        .SetDefault(-1);
426 427 428 429 430 431 432 433 434
    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);
435 436 437 438 439 440 441 442 443 444 445 446 447 448
    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({});
449
    AddComment(R"DOC(
450
CoalesceTensor Operator.
451

452
coalesce_tensor is used to make the address of Output
453 454 455 456 457 458 459 460
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
461
coalesce_tensor allows copying the value of Input to Output, or
462 463
setting the Output with a constant value, or persist the original Output
value.
464 465 466 467 468 469 470 471

)DOC");
  }
};

}  // namespace operators
}  // namespace paddle

472 473
REGISTER_OPERATOR(coalesce_tensor, paddle::operators::CoalesceTensorOp,
                  paddle::operators::CoalesceTensorOpMaker);
474
namespace ops = paddle::operators;
475
namespace plat = paddle::platform;
476
REGISTER_OP_CPU_KERNEL(
477
    coalesce_tensor,
478 479 480
    ops::CoalesceTensorOpKernel<paddle::platform::CPUDeviceContext, int>,
    ops::CoalesceTensorOpKernel<paddle::platform::CPUDeviceContext, float>,
    ops::CoalesceTensorOpKernel<paddle::platform::CPUDeviceContext, double>);
481

482
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
483
REGISTER_OP_CUDA_KERNEL(
484
    coalesce_tensor,
485 486 487 488 489
    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>);
490
#endif
491

492 493 494 495 496 497 498 499 500 501
#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 已提交
502 503 504 505 506 507 508 509 510 511
#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

512 513 514 515 516 517 518 519 520 521
#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

522 523 524 525 526 527 528 529 530
#if defined(PADDLE_WITH_MLU)
REGISTER_OP_MLU_KERNEL(
    coalesce_tensor,
    ops::CoalesceTensorOpKernel<paddle::platform::CPUDeviceContext,
                                plat::float16>,
    ops::CoalesceTensorOpKernel<paddle::platform::CPUDeviceContext, int>,
    ops::CoalesceTensorOpKernel<paddle::platform::CPUDeviceContext, float>);
#endif

531 532 533 534 535 536 537 538 539
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.",
540 541 542 543 544 545 546 547 548 549 550
            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));