wehub-resource-sync
311 lines
10 KiB
C++
311 lines
10 KiB
C++
// Copyright (c) 2023 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/utils/general_functions.h"
|
|
|
|
#include <unordered_set>
|
|
|
|
#include "paddle/common/ddim.h"
|
|
#include "paddle/common/enforce.h"
|
|
#include "paddle/common/errors.h"
|
|
|
|
#include "paddle/fluid/framework/ir/xpu/quant_utils.h"
|
|
#include "paddle/fluid/framework/scope.h"
|
|
#include "paddle/fluid/framework/variable.h"
|
|
|
|
#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
|
|
#include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
|
|
#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
|
|
#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
|
|
#include "paddle/fluid/pir/drr/src/ir_operation_factory.h"
|
|
|
|
#include "paddle/pir/include/core/builtin_op.h"
|
|
#include "paddle/pir/include/core/op_operand.h"
|
|
#include "paddle/pir/include/core/operation.h"
|
|
#include "paddle/pir/include/core/operation_utils.h"
|
|
#include "paddle/pir/include/core/program.h"
|
|
#include "paddle/pir/include/core/value.h"
|
|
#include "paddle/pir/include/pass/pass.h"
|
|
#include "paddle/pir/include/pattern_rewrite/pattern_match.h"
|
|
|
|
#include "paddle/phi/common/data_type.h"
|
|
#include "paddle/phi/common/place.h"
|
|
#include "paddle/phi/core/dense_tensor.h"
|
|
#include "paddle/phi/kernels/assign_kernel.h"
|
|
#include "paddle/phi/kernels/cast_kernel.h"
|
|
#include "paddle/phi/kernels/scale_kernel.h"
|
|
|
|
namespace pir {
|
|
|
|
void GetUsedExternalValueImpl(
|
|
std::unordered_set<Value>& defined_values, // NOLINT
|
|
std::vector<Value>& used_values, // NOLINT
|
|
const Operation& op) {
|
|
for (size_t index = 0; index < op.num_operands(); ++index) {
|
|
Value value = op.operand_source(index);
|
|
if (defined_values.find(value) == defined_values.end()) {
|
|
used_values.push_back(value);
|
|
defined_values.insert(value);
|
|
}
|
|
}
|
|
for (auto& region : op) {
|
|
for (auto& block : region) {
|
|
for (auto value : block.args()) {
|
|
defined_values.insert(value);
|
|
}
|
|
for (const auto& [_, value] : block.kwargs()) {
|
|
defined_values.insert(value);
|
|
}
|
|
}
|
|
for (auto& block : region) {
|
|
for (auto& inner_op : block) {
|
|
GetUsedExternalValueImpl(defined_values, used_values, inner_op);
|
|
}
|
|
}
|
|
}
|
|
for (size_t index = 0; index < op.num_results(); ++index) {
|
|
defined_values.insert(op.result(index));
|
|
}
|
|
}
|
|
|
|
void TensorCopySync(const phi::DenseTensor& src,
|
|
phi::DenseTensor* dst,
|
|
const phi::Place& dst_place) {
|
|
paddle::framework::TensorCopySync(src, dst_place, dst);
|
|
}
|
|
|
|
void DenseTensorCastToFp32(phi::DenseTensor* in,
|
|
phi::DenseTensor* out,
|
|
int world_size) {
|
|
auto* cpu_ctx = static_cast<phi::CPUContext*>(
|
|
phi::DeviceContextPool::Instance().Get(phi::CPUPlace()));
|
|
|
|
phi::DenseTensor fp32_tensor;
|
|
phi::DenseTensor* out_ptr = out == nullptr ? &fp32_tensor : out;
|
|
out_ptr->Resize(in->dims());
|
|
out_ptr->set_type(phi::DataType::FLOAT32);
|
|
out_ptr->set_layout(in->layout());
|
|
|
|
switch (in->dtype()) {
|
|
case phi::DataType::FLOAT16:
|
|
phi::CastKernel<phi::dtype::float16, phi::CPUContext>(
|
|
*cpu_ctx, *in, phi::DataType::FLOAT32, out_ptr);
|
|
break;
|
|
case phi::DataType::FLOAT32:
|
|
if (out == nullptr) {
|
|
if (world_size > 1) {
|
|
phi::ScaleKernel<float, phi::CPUContext>(
|
|
*cpu_ctx, *in, 1.0f / world_size, 0.f, false, in);
|
|
}
|
|
return;
|
|
} else {
|
|
phi::AssignKernel(*cpu_ctx, *in, out_ptr);
|
|
}
|
|
break;
|
|
default:
|
|
PADDLE_THROW(common::errors::InvalidType(
|
|
"Only support fp16 and fp32, but received dtype is %s.",
|
|
phi::DataTypeToString(in->dtype())));
|
|
break;
|
|
}
|
|
if (world_size > 1) {
|
|
phi::ScaleKernel<float, phi::CPUContext>(
|
|
*cpu_ctx, *out_ptr, 1.0f / world_size, 0.f, false, out_ptr);
|
|
}
|
|
if (out == nullptr) {
|
|
phi::AssignKernel(*cpu_ctx, *in, out_ptr);
|
|
}
|
|
}
|
|
|
|
Type TranslateToIrDataType(phi::DataType dtype) {
|
|
// Get Meta
|
|
IrContext* ctx = IrContext::Instance();
|
|
Type data_type = paddle::dialect::TransToIrDataType(dtype, ctx);
|
|
return data_type;
|
|
}
|
|
|
|
Operation* CreateOperationByName(const std::string& op_name,
|
|
const std::vector<Value>& inputs,
|
|
const AttributeMap& attrs,
|
|
const PatternRewriter& rewriter) {
|
|
return paddle::drr::OperationFactory::Instance().CreateOperation(
|
|
op_name, inputs, attrs, const_cast<PatternRewriter&>(rewriter));
|
|
}
|
|
|
|
Attribute CreateDataTypeAttr(IrContext* ctx, phi::DataType dtype) {
|
|
return paddle::dialect::DataTypeAttribute::get(ctx, dtype);
|
|
}
|
|
|
|
template <typename T>
|
|
T* VarGetMutable(Variable* var) {
|
|
return var->GetMutable<T>();
|
|
}
|
|
|
|
template <typename T>
|
|
bool VarIsType(Variable* var) {
|
|
return var->IsType<T>();
|
|
}
|
|
|
|
template phi::DenseTensor* VarGetMutable<phi::DenseTensor>(Variable*);
|
|
template bool VarIsType<phi::DenseTensor>(Variable*);
|
|
|
|
Variable* ScopeFindVar(Scope* scope_, const std::string& name) {
|
|
return scope_->FindVar(name);
|
|
}
|
|
|
|
Variable* ScopeGetVar(Scope* scope_, const std::string& name) {
|
|
return scope_->GetVar(name);
|
|
}
|
|
|
|
Variable* ScopeVar(Scope* scope_, const std::string& name) {
|
|
return scope_->Var(name);
|
|
}
|
|
|
|
std::vector<std::string> ScopeGetVarNames(Scope* scope_) {
|
|
return scope_->LocalVarNames();
|
|
}
|
|
|
|
Scope* GetScopeImpl(Pass* pass) {
|
|
// get scope from pass
|
|
return &pass->Get<Scope>(Pass::kParamScopeAttr);
|
|
}
|
|
|
|
std::string GetParameterNameFromValue(const Value& value) {
|
|
Operation* owner = value.defining_op();
|
|
std::string name;
|
|
if (owner->isa<ParameterOp>()) {
|
|
ParameterOp op = owner->dyn_cast<ParameterOp>();
|
|
name = op.param_name();
|
|
} else if (owner->isa<ConstantTensorOp>()) {
|
|
ConstantTensorOp op = owner->dyn_cast<ConstantTensorOp>();
|
|
name = op.tensor_name();
|
|
} else {
|
|
PADDLE_THROW(
|
|
common::errors::Unimplemented("Value must be a weight from a Parameter "
|
|
"or a ConstantTensorOp op."));
|
|
}
|
|
return name;
|
|
}
|
|
|
|
std::vector<int64_t> GetShapeFromValue(const Value& value) {
|
|
if (value.type().isa<DenseTensorType>()) {
|
|
return phi::vectorize(value.type().dyn_cast<DenseTensorType>().dims());
|
|
} else if (value.type().isa<SelectedRowsType>()) {
|
|
return phi::vectorize(value.type().dyn_cast<SelectedRowsType>().dims());
|
|
} else {
|
|
PADDLE_THROW(common::errors::InvalidArgument(
|
|
"Currently, we can only get shape for dense_tensor or selected_rows."));
|
|
}
|
|
}
|
|
|
|
Type GetDataTypeFromValue(const Value& value) {
|
|
// TODO(dev): Support other types like DenseTensor.
|
|
PADDLE_ENFORCE_EQ(value.type().isa<DenseTensorType>(),
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"Value's type must be a DenseTensorType."));
|
|
return value.type().dyn_cast<DenseTensorType>().dtype();
|
|
}
|
|
|
|
Operation* GetDefiningOpForInput(const Operation* op, uint32_t index) {
|
|
PADDLE_ENFORCE_EQ(
|
|
index < op->num_operands() && op->operand_source(index),
|
|
true,
|
|
common::errors::InvalidArgument("Input operand's index must be valid."));
|
|
return op->operand_source(index).defining_op();
|
|
}
|
|
|
|
std::vector<std::pair<Operation*, int32_t>> GetUseOpsForOutput(
|
|
const Operation* op, uint32_t index) {
|
|
PADDLE_ENFORCE_EQ(index < op->num_results(),
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"Output op result's index must be valid."));
|
|
auto result = op->result(index);
|
|
std::vector<std::pair<Operation*, int32_t>> use_ops;
|
|
for (auto it = result.use_begin(); it != result.use_end(); ++it) {
|
|
use_ops.emplace_back(it->owner(), it->index());
|
|
}
|
|
return use_ops;
|
|
}
|
|
|
|
std::vector<Value> GetUsedExternalValue(const Operation& op) {
|
|
std::unordered_set<Value> defined_values{nullptr};
|
|
std::vector<Value> used_values;
|
|
GetUsedExternalValueImpl(defined_values, used_values, op);
|
|
return used_values;
|
|
}
|
|
|
|
std::vector<Value> GetUsedExternalValue(const Block& block) {
|
|
auto& args = block.args();
|
|
std::unordered_set<Value> defined_values(args.begin(), args.end());
|
|
std::vector<Value> used_values;
|
|
for (auto& op : block) {
|
|
GetUsedExternalValueImpl(defined_values, used_values, op);
|
|
}
|
|
return used_values;
|
|
}
|
|
|
|
bool ValueIsPersistable(const Value& value) {
|
|
if (!value.defining_op()) {
|
|
return false;
|
|
}
|
|
if (value.defining_op()->num_operands() > 0) {
|
|
for (const auto& source_value : value.defining_op()->operands_source()) {
|
|
if (!ValueIsPersistable(source_value)) {
|
|
return false;
|
|
}
|
|
}
|
|
} else {
|
|
if (!value.defining_op()->isa<ParameterOp>() &&
|
|
!value.defining_op()->isa<paddle::dialect::FullOp>() &&
|
|
!value.defining_op()->isa<paddle::dialect::FullIntArrayOp>()) {
|
|
return false;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
phi::DataType GetTensorDtype(Type type) {
|
|
if (!type) {
|
|
PADDLE_THROW(
|
|
common::errors::InvalidArgument("The type of value is nullptr."));
|
|
}
|
|
if (auto dense_tensor_type = type.dyn_cast<DenseTensorType>()) {
|
|
return paddle::dialect::TransToPhiDataType(dense_tensor_type.dtype());
|
|
} else if (auto sparse_coo_tensor_type =
|
|
type.dyn_cast<paddle::dialect::SparseCooTensorType>()) {
|
|
return paddle::dialect::TransToPhiDataType(sparse_coo_tensor_type.dtype());
|
|
} else if (auto sparse_csr_tensor_type =
|
|
type.dyn_cast<paddle::dialect::SparseCsrTensorType>()) {
|
|
return paddle::dialect::TransToPhiDataType(sparse_csr_tensor_type.dtype());
|
|
} else if (auto select_rows = type.dyn_cast<SelectedRowsType>()) {
|
|
return paddle::dialect::TransToPhiDataType(select_rows.dtype());
|
|
} else if (auto dense_array =
|
|
type.dyn_cast<paddle::dialect::DenseTensorArrayType>()) {
|
|
return paddle::dialect::TransToPhiDataType(dense_array.dtype());
|
|
} else {
|
|
PADDLE_THROW(common::errors::InvalidArgument(
|
|
"Currently, we can only get phi::DataType from DenseTensorType and "
|
|
"SelectedRowsType, DenseTensorArrayType,SparseCooTensorType or "
|
|
"SparseCsrTensorType."));
|
|
}
|
|
}
|
|
|
|
phi::DataType GetValueDtype(const Value& val) {
|
|
return GetTensorDtype(val.type());
|
|
}
|
|
|
|
} // namespace pir
|