A pass which simplifies patterns of Tuple and GetTupleElement instructions in the module.
| 26 | // A pass which simplifies patterns of Tuple and GetTupleElement instructions in |
| 27 | // the module. |
| 28 | class TupleSimplifier : public HloModulePass { |
| 29 | public: |
| 30 | TupleSimplifier() : TupleSimplifier(/*exclude_entry_computation=*/false) {} |
| 31 | explicit TupleSimplifier(bool exclude_entry_computation); |
| 32 | ~TupleSimplifier() override {} |
| 33 | absl::string_view name() const override { return "tuple-simplifier"; } |
| 34 | |
| 35 | // Run tuple simplification on the given computation. Returns whether the |
| 36 | // computation was changed. |
| 37 | StatusOr<bool> Run(HloModule* module) override; |
| 38 | |
| 39 | private: |
| 40 | // When set, this pipeline stage will perform optimization of all computations |
| 41 | // apart from the module's entry computation. This is used by Graphcore's |
| 42 | // backend. |
| 43 | bool exclude_entry_computation_; |
| 44 | }; |
| 45 | |
| 46 | } // namespace xla |
| 47 |