Recompute Offload (#30233)
JZ-LIANG
4 years ago
committed by
GitHub
parent
2e80857760
commit
75936d838f
@ -0,0 +1,146 @@
|
|||||||
|
/* Copyright (c) 2020 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. */
|
||||||
|
|
||||||
|
#include "paddle/fluid/operators/memcpy_op.h"
|
||||||
|
|
||||||
|
#include <string>
|
||||||
|
|
||||||
|
namespace paddle {
|
||||||
|
namespace framework {
|
||||||
|
class OpDesc;
|
||||||
|
class Variable;
|
||||||
|
} // namespace framework
|
||||||
|
namespace imperative {
|
||||||
|
class OpBase;
|
||||||
|
} // namespace imperative
|
||||||
|
namespace platform {
|
||||||
|
struct CPUPlace;
|
||||||
|
struct CUDAPlace;
|
||||||
|
struct float16;
|
||||||
|
} // namespace platform
|
||||||
|
} // namespace paddle
|
||||||
|
|
||||||
|
namespace paddle {
|
||||||
|
namespace operators {
|
||||||
|
|
||||||
|
class MemcpyOp : public framework::OperatorWithKernel {
|
||||||
|
public:
|
||||||
|
MemcpyOp(const std::string &type, const framework::VariableNameMap &inputs,
|
||||||
|
const framework::VariableNameMap &outputs,
|
||||||
|
const framework::AttributeMap &attrs)
|
||||||
|
: OperatorWithKernel(type, inputs, outputs, attrs) {}
|
||||||
|
|
||||||
|
void InferShape(framework::InferShapeContext *ctx) const override {
|
||||||
|
auto type = ctx->GetInputsVarType("X")[0];
|
||||||
|
if (type == framework::proto::VarType::SELECTED_ROWS ||
|
||||||
|
type == framework::proto::VarType::LOD_TENSOR) {
|
||||||
|
ctx->SetOutputDim("Out", ctx->GetInputDim("X"));
|
||||||
|
if (type == framework::proto::VarType::LOD_TENSOR) {
|
||||||
|
ctx->ShareLoD("X", /*->*/ "Out");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
protected:
|
||||||
|
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());
|
||||||
|
}
|
||||||
|
|
||||||
|
framework::OpKernelType GetExpectedKernelType(
|
||||||
|
const framework::ExecutionContext &ctx) const override {
|
||||||
|
return framework::OpKernelType(
|
||||||
|
OperatorWithKernel::IndicateVarDataType(ctx, "X"),
|
||||||
|
ctx.device_context());
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
class MemcpyInferVarType : public framework::VarTypeInference {
|
||||||
|
public:
|
||||||
|
void operator()(framework::InferVarTypeContext *ctx) const override {
|
||||||
|
ctx->SyncTypeAndDataType("X", "Out");
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
class MemcpyKernel {
|
||||||
|
public:
|
||||||
|
void operator()(const framework::ExecutionContext &ctx) const {
|
||||||
|
auto *x = ctx.InputVar("X");
|
||||||
|
if (x == nullptr) {
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
PADDLE_ENFORCE_EQ(
|
||||||
|
ctx.HasOutput("Out"), true,
|
||||||
|
platform::errors::NotFound("Output(Out) of memcpy_op is not found."));
|
||||||
|
auto *out = ctx.OutputVar("Out");
|
||||||
|
platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();
|
||||||
|
auto &dev_ctx = *pool.Get(ctx.GetPlace());
|
||||||
|
auto dst_place_type = ctx.Attr<int>("dst_place_type");
|
||||||
|
framework::VisitVarType(*x, MemcpyFunctor(out, dev_ctx, dst_place_type));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
class MemcpyOpProtoMaker : public framework::OpProtoAndCheckerMaker {
|
||||||
|
public:
|
||||||
|
void Make() override {
|
||||||
|
AddInput("X", "(LoDTensor) The input variable ");
|
||||||
|
AddOutput("Out",
|
||||||
|
"(LoDTensor) The type of output "
|
||||||
|
"is the same as input X.");
|
||||||
|
AddAttr<int>("dst_place_type",
|
||||||
|
"Determine the dst place of tensor copy. "
|
||||||
|
"By Now it ONLY support CUDAPlace and CUDAPinnedPlace. Other "
|
||||||
|
"place type is Unimplemented and will cause ERROR."
|
||||||
|
"0: dst is on CPUPlace. "
|
||||||
|
"1: dst is on CUDAPlace. "
|
||||||
|
"2: dst is on CUDAPinnedPlace. "
|
||||||
|
"3: dst is on XPUPlace. ");
|
||||||
|
AddComment(R"DOC(
|
||||||
|
Memcpy Operator.
|
||||||
|
By now, it ONLY supports the memcopy between CUDAPinnedPlace and CUDAPlace,
|
||||||
|
and used as an internal op by Recompute-Offload.
|
||||||
|
You would have to update it if you want other more capacities.
|
||||||
|
|
||||||
|
Out = X, when type in [LoDTensor]
|
||||||
|
raise error if the type is not listed above.
|
||||||
|
)DOC");
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
} // namespace operators
|
||||||
|
} // namespace paddle
|
||||||
|
|
||||||
|
namespace ops = paddle::operators;
|
||||||
|
namespace plat = paddle::platform;
|
||||||
|
REGISTER_OPERATOR(
|
||||||
|
memcpy, ops::MemcpyOp, ops::MemcpyOpProtoMaker, ops::MemcpyInferVarType,
|
||||||
|
paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>,
|
||||||
|
paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>);
|
||||||
|
|
||||||
|
REGISTER_OP_CPU_KERNEL_FUNCTOR(memcpy, float, ops::MemcpyKernel, double,
|
||||||
|
ops::MemcpyKernel, int, ops::MemcpyKernel,
|
||||||
|
int64_t, ops::MemcpyKernel, bool,
|
||||||
|
ops::MemcpyKernel, plat::float16,
|
||||||
|
ops::MemcpyKernel);
|
||||||
|
|
||||||
|
#ifdef PADDLE_WITH_CUDA
|
||||||
|
REGISTER_OP_CUDA_KERNEL_FUNCTOR(memcpy, float, ops::MemcpyKernel, double,
|
||||||
|
ops::MemcpyKernel, int, ops::MemcpyKernel,
|
||||||
|
int64_t, ops::MemcpyKernel, bool,
|
||||||
|
ops::MemcpyKernel, plat::float16,
|
||||||
|
ops::MemcpyKernel);
|
||||||
|
#endif
|
||||||
@ -0,0 +1,75 @@
|
|||||||
|
/* Copyright (c) 2020 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. */
|
||||||
|
|
||||||
|
#pragma once
|
||||||
|
|
||||||
|
#include "paddle/fluid/framework/data_type.h"
|
||||||
|
#include "paddle/fluid/framework/op_registry.h"
|
||||||
|
#include "paddle/fluid/framework/var_type.h"
|
||||||
|
#include "paddle/fluid/platform/device_context.h"
|
||||||
|
|
||||||
|
namespace paddle {
|
||||||
|
namespace framework {
|
||||||
|
class LoDTensor;
|
||||||
|
class Variable;
|
||||||
|
} // namespace framework
|
||||||
|
} // namespace paddle
|
||||||
|
|
||||||
|
namespace paddle {
|
||||||
|
namespace operators {
|
||||||
|
class MemcpyFunctor {
|
||||||
|
public:
|
||||||
|
MemcpyFunctor(framework::Variable *out,
|
||||||
|
const platform::DeviceContext &dev_ctx,
|
||||||
|
const int dst_place_type)
|
||||||
|
: out_(out), dev_ctx_(dev_ctx), dst_place_type_(dst_place_type) {}
|
||||||
|
|
||||||
|
void operator()(const framework::LoDTensor &lod_tensor) const {
|
||||||
|
auto &out_tensor = *out_->GetMutable<framework::LoDTensor>();
|
||||||
|
|
||||||
|
if (dst_place_type_ == 3) {
|
||||||
|
framework::TensorCopy(lod_tensor, platform::CUDAPinnedPlace(), dev_ctx_,
|
||||||
|
&out_tensor);
|
||||||
|
} else if (dst_place_type_ == 2) {
|
||||||
|
framework::TensorCopy(lod_tensor, dev_ctx_.GetPlace(), dev_ctx_,
|
||||||
|
&out_tensor);
|
||||||
|
} else {
|
||||||
|
PADDLE_THROW(platform::errors::Unimplemented(
|
||||||
|
"memcpy dst_place_type: %d is not supported yet.", dst_place_type_));
|
||||||
|
}
|
||||||
|
out_tensor.set_lod(lod_tensor.lod());
|
||||||
|
}
|
||||||
|
|
||||||
|
void operator()(const framework::SelectedRows &rows) const {
|
||||||
|
// (JZ-LIANG) to support SelectedRows
|
||||||
|
PADDLE_THROW(platform::errors::Unimplemented(
|
||||||
|
"Memcpy for SelectedRows is NOT support yet."));
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename T>
|
||||||
|
void operator()(const T &v) const {
|
||||||
|
PADDLE_ENFORCE_EQ(
|
||||||
|
true, false,
|
||||||
|
platform::errors::PermissionDenied(
|
||||||
|
"Not support type for Memcpy op with type %s", typeid(T).name()));
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
framework::Variable *out_;
|
||||||
|
const platform::DeviceContext &dev_ctx_;
|
||||||
|
const int dst_place_type_;
|
||||||
|
};
|
||||||
|
|
||||||
|
} // namespace operators
|
||||||
|
} // namespace paddle
|
||||||
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,176 @@
|
|||||||
|
# Copyright (c) 2018 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.
|
||||||
|
|
||||||
|
from __future__ import print_function
|
||||||
|
|
||||||
|
import op_test
|
||||||
|
import numpy as np
|
||||||
|
import unittest
|
||||||
|
import paddle
|
||||||
|
import paddle.fluid.core as core
|
||||||
|
from paddle.fluid.op import Operator
|
||||||
|
import paddle.fluid as fluid
|
||||||
|
from paddle.fluid import compiler, Program, program_guard
|
||||||
|
from paddle.fluid.backward import append_backward
|
||||||
|
|
||||||
|
|
||||||
|
class TestMemcpy_FillConstant(unittest.TestCase):
|
||||||
|
def get_prog(self):
|
||||||
|
paddle.enable_static()
|
||||||
|
main_program = Program()
|
||||||
|
with program_guard(main_program):
|
||||||
|
pinned_var_name = "tensor@Pinned"
|
||||||
|
gpu_var_name = "tensor@GPU"
|
||||||
|
pinned_var = main_program.global_block().create_var(
|
||||||
|
name=pinned_var_name,
|
||||||
|
shape=[10, 10],
|
||||||
|
dtype='float32',
|
||||||
|
persistable=False,
|
||||||
|
stop_gradient=True)
|
||||||
|
gpu_var = main_program.global_block().create_var(
|
||||||
|
name=gpu_var_name,
|
||||||
|
shape=[10, 10],
|
||||||
|
dtype='float32',
|
||||||
|
persistable=False,
|
||||||
|
stop_gradient=True)
|
||||||
|
main_program.global_block().append_op(
|
||||||
|
type="fill_constant",
|
||||||
|
outputs={"Out": gpu_var_name},
|
||||||
|
attrs={
|
||||||
|
"shape": [10, 10],
|
||||||
|
"dtype": gpu_var.dtype,
|
||||||
|
"value": 1.0,
|
||||||
|
"place_type": 1
|
||||||
|
})
|
||||||
|
main_program.global_block().append_op(
|
||||||
|
type="fill_constant",
|
||||||
|
outputs={"Out": pinned_var_name},
|
||||||
|
attrs={
|
||||||
|
"shape": [10, 10],
|
||||||
|
"dtype": gpu_var.dtype,
|
||||||
|
"value": 0.0,
|
||||||
|
"place_type": 2
|
||||||
|
})
|
||||||
|
return main_program, gpu_var, pinned_var
|
||||||
|
|
||||||
|
def test_gpu_cpoy_to_pinned(self):
|
||||||
|
main_program, gpu_var, pinned_var = self.get_prog()
|
||||||
|
main_program.global_block().append_op(
|
||||||
|
type='memcpy',
|
||||||
|
inputs={'X': gpu_var},
|
||||||
|
outputs={'Out': pinned_var},
|
||||||
|
attrs={'dst_place_type': 3})
|
||||||
|
place = fluid.CUDAPlace(0)
|
||||||
|
exe = fluid.Executor(place)
|
||||||
|
gpu_, pinned_ = exe.run(main_program,
|
||||||
|
feed={},
|
||||||
|
fetch_list=[gpu_var.name, pinned_var.name])
|
||||||
|
self.assertTrue(np.allclose(gpu_, pinned_))
|
||||||
|
self.assertTrue(np.allclose(pinned_, np.ones((10, 10))))
|
||||||
|
|
||||||
|
def test_pinned_cpoy_gpu(self):
|
||||||
|
main_program, gpu_var, pinned_var = self.get_prog()
|
||||||
|
main_program.global_block().append_op(
|
||||||
|
type='memcpy',
|
||||||
|
inputs={'X': pinned_var},
|
||||||
|
outputs={'Out': gpu_var},
|
||||||
|
attrs={'dst_place_type': 2})
|
||||||
|
place = fluid.CUDAPlace(0)
|
||||||
|
exe = fluid.Executor(place)
|
||||||
|
gpu_, pinned_ = exe.run(main_program,
|
||||||
|
feed={},
|
||||||
|
fetch_list=[gpu_var.name, pinned_var.name])
|
||||||
|
self.assertTrue(np.allclose(gpu_, pinned_))
|
||||||
|
self.assertTrue(np.allclose(gpu_, np.zeros((10, 10))))
|
||||||
|
|
||||||
|
|
||||||
|
class TestMemcpyOPError(unittest.TestCase):
|
||||||
|
def get_prog(self):
|
||||||
|
paddle.enable_static()
|
||||||
|
main_program = Program()
|
||||||
|
with program_guard(main_program):
|
||||||
|
pinned_var = main_program.global_block().create_var(
|
||||||
|
name="tensor@Pinned_0",
|
||||||
|
shape=[10, 10],
|
||||||
|
dtype='float32',
|
||||||
|
persistable=False,
|
||||||
|
stop_gradient=True)
|
||||||
|
main_program.global_block().append_op(
|
||||||
|
type="fill_constant",
|
||||||
|
outputs={"Out": "tensor@Pinned_0"},
|
||||||
|
attrs={
|
||||||
|
"shape": [10, 10],
|
||||||
|
"dtype": pinned_var.dtype,
|
||||||
|
"value": 0.0,
|
||||||
|
"place_type": 2
|
||||||
|
})
|
||||||
|
return main_program, pinned_var
|
||||||
|
|
||||||
|
def test_SELECTED_ROWS(self):
|
||||||
|
main_program, pinned_var = self.get_prog()
|
||||||
|
selected_row_var = main_program.global_block().create_var( \
|
||||||
|
name="selected_row_0", dtype="float32", persistable=False, \
|
||||||
|
type=fluid.core.VarDesc.VarType.SELECTED_ROWS, stop_gradient=True)
|
||||||
|
main_program.global_block().append_op(
|
||||||
|
type="fill_constant",
|
||||||
|
outputs={"Out": selected_row_var},
|
||||||
|
attrs={
|
||||||
|
"shape": selected_row_var.shape,
|
||||||
|
"dtype": selected_row_var.dtype,
|
||||||
|
"value": 1.0,
|
||||||
|
"place_type": 1
|
||||||
|
})
|
||||||
|
main_program.global_block().append_op(
|
||||||
|
type='memcpy',
|
||||||
|
inputs={'X': selected_row_var},
|
||||||
|
outputs={'Out': pinned_var},
|
||||||
|
attrs={'dst_place_type': 3})
|
||||||
|
with self.assertRaises(NotImplementedError):
|
||||||
|
place = fluid.CUDAPlace(0)
|
||||||
|
exe = fluid.Executor(place)
|
||||||
|
selected_row_var_, pinned_ = exe.run(
|
||||||
|
main_program,
|
||||||
|
feed={},
|
||||||
|
fetch_list=[selected_row_var.name, pinned_var.name])
|
||||||
|
|
||||||
|
def test_OTHER_PLACE_NotImplementedError(self):
|
||||||
|
main_program, pinned_var = self.get_prog()
|
||||||
|
lod_tensor_var = main_program.global_block().create_var( \
|
||||||
|
name="lod_tensor_0", dtype="float32", persistable=False, stop_gradient=True)
|
||||||
|
main_program.global_block().append_op(
|
||||||
|
type="fill_constant",
|
||||||
|
outputs={"Out": lod_tensor_var},
|
||||||
|
attrs={
|
||||||
|
"shape": lod_tensor_var.shape,
|
||||||
|
"dtype": lod_tensor_var.dtype,
|
||||||
|
"value": 1.0,
|
||||||
|
"place_type": 0
|
||||||
|
})
|
||||||
|
main_program.global_block().append_op(
|
||||||
|
type='memcpy',
|
||||||
|
inputs={'X': pinned_var},
|
||||||
|
outputs={'Out': lod_tensor_var},
|
||||||
|
attrs={'dst_place_type': 0, })
|
||||||
|
with self.assertRaises(NotImplementedError):
|
||||||
|
place = fluid.CUDAPlace(0)
|
||||||
|
exe = fluid.Executor(place)
|
||||||
|
lod_tensor_var_, pinned_ = exe.run(
|
||||||
|
main_program,
|
||||||
|
feed={},
|
||||||
|
fetch_list=[lod_tensor_var.name, pinned_var.name])
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
paddle.enable_static()
|
||||||
|
unittest.main()
|
||||||
Loading…
Reference in new issue