wehub-resource-sync
75 lines
2.6 KiB
C++
75 lines
2.6 KiB
C++
// Copyright (c) 2022 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/pir/dialect/operator/ir/pd_api.h"
|
|
#include "paddle/fluid/primitive/base/lazy_tensor.h"
|
|
#include "paddle/fluid/primitive/decomp_utils/decomp_utils.h"
|
|
|
|
namespace paddle::primitive {
|
|
template <>
|
|
void set_output<LazyTensor>(const Tensor& x_tmp, Tensor* x) {
|
|
x->set_impl(x_tmp.impl());
|
|
}
|
|
|
|
template <>
|
|
void by_pass<LazyTensor>(const Tensor& x, Tensor* real_out) {
|
|
pir::Value x_res = std::static_pointer_cast<LazyTensor>(x.impl())->value();
|
|
auto op_res = paddle::dialect::assign(x_res);
|
|
Tensor out(std::make_shared<LazyTensor>(op_res));
|
|
set_output<LazyTensor>(out, real_out);
|
|
}
|
|
|
|
/**
|
|
* @brief set output with empty grads in pir.
|
|
*
|
|
* In pir, we use None type to express
|
|
* that value is not available.
|
|
* Some outputs in vjp are marked as unnecessary
|
|
* by stop_gradient with True. Therefore the
|
|
* type of those outputs that are unnecessary will
|
|
* be set with None.
|
|
*
|
|
*/
|
|
void SetEmptyGrad(const std::vector<std::vector<Tensor>>& outputs,
|
|
const std::vector<std::vector<bool>>& stop_gradients) {
|
|
for (size_t i = 0; i < outputs.size(); ++i) {
|
|
for (size_t j = 0; j < outputs[i].size(); ++j) {
|
|
if (stop_gradients[i][j] && outputs[i][j].impl()) {
|
|
std::static_pointer_cast<primitive::LazyTensor>(outputs[i][j].impl())
|
|
->set_empty();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
std::vector<std::vector<Tensor>> ConstructVjpResultByStopGradients(
|
|
const std::vector<std::vector<Tensor>>& outputs,
|
|
const std::vector<std::vector<bool>>& stop_gradients) {
|
|
SetEmptyGrad(outputs, stop_gradients);
|
|
std::vector<std::vector<Tensor>> vjp_results(outputs.size());
|
|
for (size_t i = 0; i < outputs.size(); ++i) {
|
|
vjp_results[i].reserve(outputs[i].size());
|
|
for (size_t j = 0; j < outputs[i].size(); ++j) {
|
|
if (stop_gradients[i][j]) {
|
|
// Use Tensor's impl is nullptr to indicate it has no gradient
|
|
vjp_results[i].emplace_back(Tensor());
|
|
} else {
|
|
vjp_results[i].emplace_back(outputs[i][j]);
|
|
}
|
|
}
|
|
}
|
|
return vjp_results;
|
|
}
|
|
|
|
} // namespace paddle::primitive
|