================ @@ -0,0 +1,184 @@ +//===----------------------------------------------------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include "PassDetail.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/IR/Region.h" +#include "mlir/Support/LogicalResult.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "clang/CIR/Dialect/IR/CIRDialect.h" +#include "clang/CIR/Dialect/Passes.h" +#include "llvm/ADT/SmallVector.h" + +using namespace mlir; +using namespace cir; + +//===----------------------------------------------------------------------===// +// Rewrite patterns +//===----------------------------------------------------------------------===// + +namespace { + +/// Simplify suitable ternary operations into select operations. +/// +/// For now we only simplify those ternary operations whose true and false +/// branches directly yield a value or a constant. That is, both of the true and +/// the false branch must either contain a cir.yield operation as the only +/// operation in the branch, or contain a cir.const operation followed by a +/// cir.yield operation that yields the constant value. +/// +/// For example, we will simplify the following ternary operation: +/// +/// %0 = cir.ternary (%condition, true { +/// %1 = cir.const ... +/// cir.yield %1 +/// } false { +/// cir.yield %2 +/// }) +/// +/// into the following sequence of operations: +/// +/// %1 = cir.const ... +/// %0 = cir.select if %condition then %1 else %2 +struct SimplifyTernary final : public OpRewritePattern<TernaryOp> { + using OpRewritePattern<TernaryOp>::OpRewritePattern; + + LogicalResult matchAndRewrite(TernaryOp op, + PatternRewriter &rewriter) const override { + if (op->getNumResults() != 1) + return mlir::failure(); + + if (!isSimpleTernaryBranch(op.getTrueRegion()) || + !isSimpleTernaryBranch(op.getFalseRegion())) + return mlir::failure(); + + cir::YieldOp trueBranchYieldOp = + mlir::cast<cir::YieldOp>(op.getTrueRegion().front().getTerminator()); + cir::YieldOp falseBranchYieldOp = + mlir::cast<cir::YieldOp>(op.getFalseRegion().front().getTerminator()); + mlir::Value trueValue = trueBranchYieldOp.getArgs()[0]; + mlir::Value falseValue = falseBranchYieldOp.getArgs()[0]; + + rewriter.inlineBlockBefore(&op.getTrueRegion().front(), op); + rewriter.inlineBlockBefore(&op.getFalseRegion().front(), op); + rewriter.eraseOp(trueBranchYieldOp); + rewriter.eraseOp(falseBranchYieldOp); + rewriter.replaceOpWithNewOp<cir::SelectOp>(op, op.getCond(), trueValue, + falseValue); + + return mlir::success(); + } + +private: + bool isSimpleTernaryBranch(mlir::Region ®ion) const { + if (!region.hasOneBlock()) + return false; + + mlir::Block &onlyBlock = region.front(); + mlir::Block::OpListType &ops = onlyBlock.getOperations(); + + // The region/block could only contain at most 2 operations. + if (ops.size() > 2) + return false; + + if (ops.size() == 1) { + // The region/block only contain a cir.yield operation. + return true; + } + + // Check whether the region/block contains a cir.const followed by a + // cir.yield that yields the value. + auto yieldOp = mlir::cast<cir::YieldOp>(onlyBlock.getTerminator()); + auto yieldValueDefOp = mlir::dyn_cast_if_present<cir::ConstantOp>( + yieldOp.getArgs()[0].getDefiningOp()); + return yieldValueDefOp && yieldValueDefOp->getBlock() == &onlyBlock; + } +}; + +struct SimplifySelect : public OpRewritePattern<SelectOp> { ---------------- andykaylor wrote:
It would be great if this had an explanatory comment like `SimplifyTernary` does. If there are going to be more cases added later, a general comment is fine. Otherwise, the two cases handled here can be explained pretty easily here. https://github.com/llvm/llvm-project/pull/138317 _______________________________________________ cfe-commits mailing list cfe-commits@lists.llvm.org https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits
