|
| 1 | +// Copyright 2024 The XLS Authors |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, software |
| 10 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +// See the License for the specific language governing permissions and |
| 13 | +// limitations under the License. |
| 14 | + |
| 15 | +#include <cassert> |
| 16 | +#include <optional> |
| 17 | +#include <utility> |
| 18 | + |
| 19 | +#include "llvm/include/llvm/ADT/SmallVector.h" |
| 20 | +#include "mlir/include/mlir/IR/BuiltinOps.h" |
| 21 | +#include "mlir/include/mlir/IR/Value.h" |
| 22 | +#include "mlir/include/mlir/IR/Visitors.h" |
| 23 | +#include "mlir/include/mlir/Pass/Pass.h" // IWYU pragma: keep |
| 24 | +#include "mlir/include/mlir/Support/LLVM.h" |
| 25 | +#include "xls/contrib/mlir/IR/xls_ops.h" |
| 26 | +#include "xls/contrib/mlir/transforms/passes.h" // IWYU pragma: keep |
| 27 | + |
| 28 | +namespace mlir::xls { |
| 29 | + |
| 30 | +#define GEN_PASS_DEF_SKIPEMPTYTOPEPROCPASS |
| 31 | +#include "xls/contrib/mlir/transforms/passes.h.inc" |
| 32 | + |
| 33 | +namespace { |
| 34 | + |
| 35 | +bool IsEmptyEproc(EprocOp eproc) { |
| 36 | + return eproc.getBody().front().without_terminator().empty(); |
| 37 | +} |
| 38 | + |
| 39 | +class SkipEmptyTopEprocPass |
| 40 | + : public impl::SkipEmptyTopEprocPassBase<SkipEmptyTopEprocPass> { |
| 41 | + using SkipEmptyTopEprocPassBase::SkipEmptyTopEprocPassBase; |
| 42 | + |
| 43 | + public: |
| 44 | + void runOnOperation() override { |
| 45 | + if (!top_proc_name.has_value()) { |
| 46 | + return; |
| 47 | + } |
| 48 | + |
| 49 | + ModuleOp module_op = getOperation(); |
| 50 | + auto eprocs = llvm::to_vector(module_op.getOps<EprocOp>()); |
| 51 | + if (eprocs.size() != 2) { |
| 52 | + return; |
| 53 | + } |
| 54 | + |
| 55 | + EprocOp new_top_eproc = eprocs[0]; |
| 56 | + EprocOp current_empty_top_eproc = eprocs[1]; |
| 57 | + assert(!(new_top_eproc.getSymName() == *top_proc_name && |
| 58 | + current_empty_top_eproc.getSymName() == *top_proc_name) && |
| 59 | + "Only one eproc must have the top-level name"); |
| 60 | + if (new_top_eproc.getSymName() == *top_proc_name) { |
| 61 | + std::swap(new_top_eproc, current_empty_top_eproc); |
| 62 | + } else if (current_empty_top_eproc.getSymName() != *top_proc_name) { |
| 63 | + // Nothing to do here then. |
| 64 | + return; |
| 65 | + } |
| 66 | + |
| 67 | + if (!IsEmptyEproc(current_empty_top_eproc)) { |
| 68 | + // Nothing to do here. This top proc is not empty. |
| 69 | + return; |
| 70 | + } |
| 71 | + |
| 72 | + // Step 1: Erase all instances of this empty eproc. |
| 73 | + auto uses = *current_empty_top_eproc.getSymbolUses(module_op); |
| 74 | + for (auto& use : uses) { |
| 75 | + use.getUser()->erase(); |
| 76 | + } |
| 77 | + |
| 78 | + // Step 2: Change the name of the non-empty non-top eproc to be the name of |
| 79 | + // the empty one. Both eproc and instances need to be updated. |
| 80 | + uses = *new_top_eproc.getSymbolUses(module_op); |
| 81 | + for (auto& use : uses) { |
| 82 | + cast<InstantiateEprocOp>(use.getUser()) |
| 83 | + .setEproc(current_empty_top_eproc.getSymName()); |
| 84 | + } |
| 85 | + new_top_eproc.setSymName(current_empty_top_eproc.getSymName()); |
| 86 | + |
| 87 | + // Step 3: Erase the empty eproc. |
| 88 | + current_empty_top_eproc.erase(); |
| 89 | + } |
| 90 | +}; |
| 91 | + |
| 92 | +} // namespace |
| 93 | +} // namespace mlir::xls |
0 commit comments