Files
paddlepaddle--paddle/paddle/pir/include/core/operation_utils.h
T
2026-07-13 12:40:42 +08:00

146 lines
4.9 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.
#pragma once
#include <initializer_list>
#include <memory>
#include "paddle/pir/include/core/attribute.h"
#include "paddle/pir/include/core/dll_decl.h"
#include "paddle/pir/include/core/op_info.h"
#include "paddle/pir/include/core/op_result.h"
#include "paddle/pir/include/core/region.h"
#include "paddle/pir/include/core/type.h"
#include "paddle/pir/include/core/value.h"
namespace pir {
class Block;
using AttributeMap = std::unordered_map<std::string, Attribute>;
using PropertyMap = std::unordered_map<std::string, Property>;
//===----------------------------------------------------------------------===//
// OperationArgument
//===----------------------------------------------------------------------===//
// This represents an operation arguments in an combined form, suitable for use
// with the builder APIs.
struct OperationArgument {
std::vector<Value> inputs;
AttributeMap attributes;
std::vector<Type> output_types;
OpInfo info;
std::vector<Block*> successors;
std::vector<std::unique_ptr<Region>> regions;
public:
PADDLE_API OperationArgument(IrContext* ir_context, const std::string& name);
explicit OperationArgument(OpInfo info) : info(info) {}
OperationArgument(const std::vector<Value>& inputs,
const AttributeMap& attributes,
const std::vector<Type>& types,
OpInfo info,
const std::vector<Block*> successors = {})
: inputs(inputs),
attributes(attributes),
output_types(types),
info(info),
successors(successors) {}
OperationArgument& operator=(const OperationArgument&) = delete;
OperationArgument(const OperationArgument&) = delete;
void AddInput(Value input) { inputs.emplace_back(input); }
template <class InputIt>
void AddInputs(InputIt first, InputIt last);
void AddInputs(std::initializer_list<Value> value_list) {
AddInputs(std::begin(value_list), std::end(value_list));
}
template <class ValueContainer>
void AddInputs(const ValueContainer& value_container) {
AddInputs(std::begin(value_container), std::end(value_container));
}
/// Add Output.
void AddOutput(Type type) { output_types.emplace_back(type); }
template <class InputIt>
void AddOutputs(InputIt first, InputIt last);
void AddOutputs(std::initializer_list<Type> type_list) {
AddOutputs(std::begin(type_list), std::end(type_list));
}
template <class TypeContainer>
void AddOutputs(const TypeContainer& type_container) {
AddOutputs(std::begin(type_container), std::end(type_container));
}
/// Add an attribute with the specified name.
void AddAttribute(const std::string& name, Attribute attr) {
attributes[name] = attr;
}
/// Add an array of named attributes.
template <class InputIt>
void AddAttributes(InputIt first, InputIt last);
template <class AttrContainer>
void AddAttributes(const AttrContainer& attr_container) {
AddAttributes(std::begin(attr_container), std::end(attr_container));
}
/// Get the context held by this operation state.
IrContext* getContext() const { return info.ir_context(); }
void AddSuccessor(Block* successor) { successors.emplace_back(successor); }
/// Create a region that should be attached to the operation. These regions
/// can be filled in immediately without waiting for Operation to be
/// created. When it is, the region bodies will be transferred.
Region& AddRegion();
/// Take a region that should be attached to the Operation. The body of the
/// region will be transferred when the Operation is created. If the
/// region is nullptr, a new empty region will be attached to the Operation.
PADDLE_API void AddRegion(std::unique_ptr<Region>&& region);
// This interface is equivalent to calling AddRegion(nullptr) 'size' times.
void AddRegions(size_t size);
};
template <class InputIt>
void OperationArgument::AddInputs(InputIt first, InputIt last) {
while (first != last) {
AddInput(*first++);
}
}
template <class InputIt>
void OperationArgument::AddOutputs(InputIt first, InputIt last) {
while (first != last) {
output_types.emplace_back(*first++);
}
}
template <class InputIt>
void OperationArgument::AddAttributes(InputIt first, InputIt last) {
while (first != last) {
attributes[first->first] = first->second;
++first;
}
}
} // namespace pir