wehub-resource-sync
89 lines
2.8 KiB
C++
89 lines
2.8 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 "paddle/pir/include/core/op_base.h"
|
|
|
|
namespace pir {
|
|
|
|
class TuplePushOp;
|
|
class TuplePopOp;
|
|
///
|
|
/// \brief This interface marks the op can create a container.
|
|
///
|
|
class ContainerOpInterface : public OpInterfaceBase<ContainerOpInterface> {
|
|
public:
|
|
struct Concept {
|
|
Value (*container_)(Operation* op);
|
|
Value (*inlet_)(Operation* op);
|
|
Value (*outlet_)(Operation* op);
|
|
size_t (*tuple_size_)(Operation* op);
|
|
Value (*inlet_element_)(Operation* op, size_t index);
|
|
Value (*outlet_element_)(Operation* op, size_t index);
|
|
};
|
|
|
|
template <class ConcreteOp>
|
|
struct Model : public Concept {
|
|
Model()
|
|
: Concept{container,
|
|
inlet,
|
|
outlet,
|
|
tuple_size,
|
|
inlet_element,
|
|
inlet_element} {}
|
|
static Value container(Operation* op) {
|
|
return op->dyn_cast<ConcreteOp>().container();
|
|
}
|
|
static Value inlet(Operation* op) {
|
|
return op->dyn_cast<ConcreteOp>().inlet();
|
|
}
|
|
static Value outlet(Operation* op) {
|
|
return op->dyn_cast<ConcreteOp>().outlet();
|
|
}
|
|
static size_t tuple_size(Operation* op) {
|
|
return op->dyn_cast<ConcreteOp>().tuple_size();
|
|
}
|
|
static Value inlet_element(Operation* op, size_t index) {
|
|
return op->dyn_cast<ConcreteOp>().container();
|
|
}
|
|
static Value outlet_element(Operation* op, size_t index) {
|
|
return op->dyn_cast<ConcreteOp>().container();
|
|
}
|
|
};
|
|
|
|
Value container() { return impl_->container_(operation()); }
|
|
Value inlet() { return impl_->inlet_(operation()); }
|
|
Value outlet() { return impl_->outlet_(operation()); }
|
|
size_t tuple_size() { return impl_->tuple_size_(operation()); }
|
|
Value inlet_element(size_t index) {
|
|
return impl_->inlet_element_(operation(), index);
|
|
}
|
|
Value outlet_element(size_t index) {
|
|
return impl_->outlet_element_(operation(), index);
|
|
}
|
|
|
|
TuplePushOp tuple_push_op();
|
|
TuplePopOp tuple_pop_op();
|
|
/// Constructor
|
|
ContainerOpInterface(const pir::Operation* op, Concept* impl)
|
|
: OpInterfaceBase<ContainerOpInterface>(op), impl_(impl) {}
|
|
|
|
private:
|
|
Concept* impl_;
|
|
};
|
|
} // namespace pir
|
|
|
|
IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::ContainerOpInterface)
|