Changeset View
Changeset View
Standalone View
Standalone View
mlir/include/mlir/IR/OpConnectivity.h
- This file was added.
| //===- OpConnectivity.h -----------------------------------------*- C++ -*-===// | |||||
| // | |||||
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. | |||||
| // See https://llvm.org/LICENSE.txt for license information. | |||||
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | |||||
| // | |||||
| //===----------------------------------------------------------------------===// | |||||
| // Utility functions to get the prodcuers and consumer operations of a given | |||||
| // operation | |||||
| //===----------------------------------------------------------------------===// | |||||
| #ifndef MLIR_IR_CONNECTIVITY_H | |||||
| #define MLIR_IR_CONNECTIVITY_H | |||||
| #include "mlir/Dialect/Quant/QuantOps.h" | |||||
| #include "mlir/Dialect/Quant/QuantTypes.h" | |||||
| #include "mlir/IR/AffineMap.h" // from @llvm-project | |||||
| #include "mlir/IR/Attributes.h" // from @llvm-project | |||||
| #include "mlir/IR/BlockAndValueMapping.h" // from @llvm-project | |||||
| #include "mlir/IR/Builders.h" // from @llvm-project | |||||
| #include "mlir/IR/Location.h" // from @llvm-project | |||||
| #include "mlir/IR/MLIRContext.h" // from @llvm-project | |||||
| #include "mlir/IR/Operation.h" // from @llvm-project | |||||
| #include "mlir/IR/PatternMatch.h" // from @llvm-project | |||||
| #include "mlir/Interfaces/InferTypeOpInterface.h" // from @llvm-project | |||||
| #include "mlir/Pass/Pass.h" // from @llvm-project | |||||
| #include "mlir/Transforms/DialectConversion.h" // from @llvm-project | |||||
| #include <iostream> | |||||
| #include <list> | |||||
| #include <set> | |||||
| #include <vector> | |||||
| using namespace std; | |||||
| namespace mlir { | |||||
| /* | |||||
| 0. Consumer/Producer/Block block index | |||||
| 1. Operation block index | |||||
| 2. Consumer/Producer/Block slot index | |||||
| 3. Operation slot index | |||||
| 4. Operation consuming/generating the value | |||||
| 5. Consumer/Producer operation generating/consuming Value | |||||
| nullptr if Value is generated by a BlockAttribute | |||||
| 6. Block generating Value | |||||
| nullptr if value is produced by an Op | |||||
| 7. Value can be a OpResult or a Operand | |||||
| 8. direction of connection | |||||
| */ | |||||
| enum connection_direction { operation_2_consumer, producer_2_operation }; | |||||
| enum producer_type { block_op, operation_op }; | |||||
| struct ArgumentInfo { | |||||
| bool block_or_op; | |||||
| Block *external_block; | |||||
| Operation *external_op; | |||||
| unsigned int external_block_op_index; | |||||
| unsigned int block_index; | |||||
| }; | |||||
| class connection { | |||||
| public: | |||||
| connection() {} | |||||
| connection(unsigned external_slot_index, unsigned operation_slot_index, | |||||
| Operation *operation, Operation *external_operation, | |||||
| Block *block_external_operation, Value value, | |||||
| connection_direction direction) { | |||||
| this->external_slot_index = external_slot_index; | |||||
| this->operation_slot_index = operation_slot_index; | |||||
| this->operation = operation; | |||||
| this->external_operation = external_operation; | |||||
| this->block_external_operation = block_external_operation; | |||||
| this->value = value; | |||||
| this->direction = direction; | |||||
| } | |||||
| unsigned external_slot_index; | |||||
| unsigned operation_slot_index; | |||||
| Operation *operation; | |||||
| Operation *external_operation; | |||||
| Block *block_external_operation; | |||||
| Value value; | |||||
| connection_direction direction; | |||||
| producer_type is_producer_block_or_operation() { | |||||
| if (this->external_operation) | |||||
| return operation_op; | |||||
| else if (this->block_external_operation) | |||||
| return block_op; | |||||
| } | |||||
| }; | |||||
| std::vector<connection> ExtractConsumers(Value v); | |||||
| std::vector<connection> ExtractConsumers(Operation *i); | |||||
| std::vector<connection> ExtractProducers(Operation *i); | |||||
| std::list<Operation *> ExtractUniqueConsumers_Operations(Operation *i); | |||||
| std::list<Operation *> ExtractUniqueProducers_Operations(Operation *i); | |||||
| template <class T> | |||||
| inline std::vector<connection> ExtractConsumers_ofType(Operation *i) { | |||||
| std::vector<connection> conn = ExtractConsumers(i); | |||||
| std::vector<connection> conn_filtered; | |||||
| for (auto c : conn) { | |||||
| if (llvm::isa<T>(c.external_operation)) | |||||
| conn_filtered.push_back(c); | |||||
| } | |||||
| return conn_filtered; | |||||
| } | |||||
| template <class T> | |||||
| inline std::vector<connection> ExtractProducers_ofType(Operation *i) { | |||||
| std::vector<connection> prod = ExtractProducers(i); | |||||
| std::vector<connection> prod_filtered; | |||||
| for (auto p : prod) { | |||||
| if (p.is_producer_block_or_operation() == operation_op) { | |||||
| if (llvm::isa<T>(p.external_operation)) | |||||
| prod_filtered.push_back(p); | |||||
| } | |||||
| } | |||||
| return prod_filtered; | |||||
| } | |||||
| SmallVector<int64_t, 4> arrayref_2_smallvector(ArrayRef<int64_t> s); | |||||
| } // namespace mlir | |||||
| #endif // MLIR_IR_CONNECTIVITY_H | |||||