Skip to content

Commit d8968f9

Browse files
AdUhTkJmAdUhTkJm
committed
[CIR][ThroughMLIR] Lower simple SwitchOp
1 parent c6a4f77 commit d8968f9

File tree

5 files changed

+286
-19
lines changed

5 files changed

+286
-19
lines changed

‎clang/include/clang/CIR/Passes.h

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,11 @@
1818
#include <memory>
1919

2020
namespace cir {
21+
/// Create a pass for transforming CIR operations to more 'scf' dialect-friendly
22+
/// forms. It rewrites operations that aren't supported by 'scf', such as breaks
23+
/// and continues.
24+
std::unique_ptr<mlir::Pass> createMLIRCoreDialectsLoweringPreparePass();
25+
2126
/// Create a pass for lowering from MLIR builtin dialects such as `Affine` and
2227
/// `Std`, to the LLVM dialect for codegen.
2328
std::unique_ptr<mlir::Pass> createConvertMLIRToLLVMPass();

‎clang/lib/CIR/Lowering/ThroughMLIR/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ add_clang_library(clangCIRLoweringThroughMLIR
99
LowerCIRLoopToSCF.cpp
1010
LowerCIRToMLIR.cpp
1111
LowerMLIRToLLVM.cpp
12+
MLIRCoreDialectsLoweringPrepare.cpp
1213

1314
DEPENDS
1415
MLIRCIROpsIncGen

‎clang/lib/CIR/Lowering/ThroughMLIR/LowerCIRToMLIR.cpp

Lines changed: 110 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1496,29 +1496,118 @@ class CIRTrapOpLowering : public mlir::OpConversionPattern<cir::TrapOp> {
14961496
}
14971497
};
14981498

1499+
class CIRSwitchOpLowering : public mlir::OpConversionPattern<cir::SwitchOp> {
1500+
public:
1501+
using OpConversionPattern<cir::SwitchOp>::OpConversionPattern;
1502+
1503+
mlir::LogicalResult
1504+
matchAndRewrite(cir::SwitchOp op, OpAdaptor adaptor,
1505+
mlir::ConversionPatternRewriter &rewriter) const override {
1506+
rewriter.setInsertionPointAfter(op);
1507+
llvm::SmallVector<CaseOp> cases;
1508+
if (!op.isSimpleForm(cases))
1509+
mlir::emitError(op.getLoc(), "not yet implemented");
1510+
1511+
llvm::SmallVector<int64_t> caseValues;
1512+
// Maps the index of a CaseOp in `cases`, to the index in `caseValues`.
1513+
// This is necessary because some CaseOp might carry 0 or multiple values.
1514+
llvm::DenseMap<size_t, unsigned> indexMap;
1515+
caseValues.reserve(cases.size());
1516+
for (auto [i, caseOp] : llvm::enumerate(cases)) {
1517+
switch (caseOp.getKind()) {
1518+
case CaseOpKind::Equal: {
1519+
auto valueAttr = caseOp.getValue()[0];
1520+
auto value = cast<cir::IntAttr>(valueAttr);
1521+
indexMap[i] = caseValues.size();
1522+
caseValues.push_back(value.getUInt());
1523+
break;
1524+
}
1525+
case CaseOpKind::Default:
1526+
break;
1527+
case CaseOpKind::Range:
1528+
case CaseOpKind::Anyof:
1529+
mlir::emitError(op.getLoc(), "not yet implemented");
1530+
}
1531+
}
1532+
1533+
auto operand = adaptor.getOperands()[0];
1534+
// `scf.index_switch` expects an index of type `index`.
1535+
auto indexType = mlir::IndexType::get(getContext());
1536+
auto indexCast = rewriter.create<mlir::arith::IndexCastOp>(
1537+
op.getLoc(), indexType, operand);
1538+
auto indexSwitch = rewriter.create<mlir::scf::IndexSwitchOp>(
1539+
op.getLoc(), mlir::TypeRange{}, indexCast, caseValues, cases.size());
1540+
1541+
bool metDefault = false;
1542+
for (auto [i, caseOp] : llvm::enumerate(cases)) {
1543+
auto &region = caseOp.getRegion();
1544+
switch (caseOp.getKind()) {
1545+
case CaseOpKind::Equal: {
1546+
auto &caseRegion = indexSwitch.getCaseRegions()[indexMap[i]];
1547+
rewriter.inlineRegionBefore(region, caseRegion, caseRegion.end());
1548+
break;
1549+
}
1550+
case CaseOpKind::Default: {
1551+
auto &defaultRegion = indexSwitch.getDefaultRegion();
1552+
rewriter.inlineRegionBefore(region, defaultRegion, defaultRegion.end());
1553+
metDefault = true;
1554+
break;
1555+
}
1556+
case CaseOpKind::Range:
1557+
case CaseOpKind::Anyof:
1558+
llvm_unreachable("NYI");
1559+
}
1560+
}
1561+
1562+
// `scf.index_switch` expects its default region to contain exactly one
1563+
// block. If we don't have a default region in `cir.switch`, we need to
1564+
// supply it here.
1565+
if (!metDefault) {
1566+
auto &defaultRegion = indexSwitch.getDefaultRegion();
1567+
mlir::Block *block =
1568+
rewriter.createBlock(&defaultRegion, defaultRegion.end());
1569+
rewriter.setInsertionPointToEnd(block);
1570+
rewriter.create<mlir::scf::YieldOp>(op.getLoc());
1571+
}
1572+
1573+
// The final `cir.break` should be replaced to `scf.yield`.
1574+
// After MLIRLoweringPrepare pass, every case must end with a `cir.break`.
1575+
for (auto &region : indexSwitch.getCaseRegions()) {
1576+
auto &lastBlock = region.back();
1577+
auto &lastOp = lastBlock.back();
1578+
assert(isa<BreakOp>(lastOp));
1579+
rewriter.setInsertionPointAfter(&lastOp);
1580+
rewriter.replaceOpWithNewOp<mlir::scf::YieldOp>(&lastOp);
1581+
}
1582+
1583+
rewriter.replaceOp(op, indexSwitch);
1584+
1585+
return mlir::success();
1586+
}
1587+
};
1588+
14991589
void populateCIRToMLIRConversionPatterns(mlir::RewritePatternSet &patterns,
15001590
mlir::TypeConverter &converter) {
15011591
patterns.add<CIRReturnLowering, CIRBrOpLowering>(patterns.getContext());
15021592

1503-
patterns
1504-
.add<CIRATanOpLowering, CIRCmpOpLowering, CIRCallOpLowering,
1505-
CIRUnaryOpLowering, CIRBinOpLowering, CIRLoadOpLowering,
1506-
CIRConstantOpLowering, CIRStoreOpLowering, CIRAllocaOpLowering,
1507-
CIRFuncOpLowering, CIRScopeOpLowering, CIRBrCondOpLowering,
1508-
CIRTernaryOpLowering, CIRYieldOpLowering, CIRCosOpLowering,
1509-
CIRGlobalOpLowering, CIRGetGlobalOpLowering, CIRCastOpLowering,
1510-
CIRPtrStrideOpLowering, CIRGetElementOpLowering, CIRSqrtOpLowering,
1511-
CIRCeilOpLowering, CIRExp2OpLowering, CIRExpOpLowering,
1512-
CIRFAbsOpLowering, CIRAbsOpLowering, CIRFloorOpLowering,
1513-
CIRLog10OpLowering, CIRLog2OpLowering, CIRLogOpLowering,
1514-
CIRRoundOpLowering, CIRPtrStrideOpLowering, CIRSinOpLowering,
1515-
CIRShiftOpLowering, CIRBitClzOpLowering, CIRBitCtzOpLowering,
1516-
CIRBitPopcountOpLowering, CIRBitClrsbOpLowering, CIRBitFfsOpLowering,
1517-
CIRBitParityOpLowering, CIRIfOpLowering, CIRVectorCreateLowering,
1518-
CIRVectorInsertLowering, CIRVectorExtractLowering,
1519-
CIRVectorCmpOpLowering, CIRACosOpLowering, CIRASinOpLowering,
1520-
CIRUnreachableOpLowering, CIRTanOpLowering, CIRTrapOpLowering>(
1521-
converter, patterns.getContext());
1593+
patterns.add<
1594+
CIRSwitchOpLowering, CIRGetElementOpLowering, CIRATanOpLowering,
1595+
CIRCmpOpLowering, CIRCallOpLowering, CIRUnaryOpLowering, CIRBinOpLowering,
1596+
CIRLoadOpLowering, CIRConstantOpLowering, CIRStoreOpLowering,
1597+
CIRAllocaOpLowering, CIRFuncOpLowering, CIRScopeOpLowering,
1598+
CIRBrCondOpLowering, CIRTernaryOpLowering, CIRYieldOpLowering,
1599+
CIRCosOpLowering, CIRGlobalOpLowering, CIRGetGlobalOpLowering,
1600+
CIRCastOpLowering, CIRPtrStrideOpLowering, CIRSqrtOpLowering,
1601+
CIRCeilOpLowering, CIRExp2OpLowering, CIRExpOpLowering, CIRFAbsOpLowering,
1602+
CIRAbsOpLowering, CIRFloorOpLowering, CIRLog10OpLowering,
1603+
CIRLog2OpLowering, CIRLogOpLowering, CIRRoundOpLowering,
1604+
CIRPtrStrideOpLowering, CIRSinOpLowering, CIRShiftOpLowering,
1605+
CIRBitClzOpLowering, CIRBitCtzOpLowering, CIRBitPopcountOpLowering,
1606+
CIRBitClrsbOpLowering, CIRBitFfsOpLowering, CIRBitParityOpLowering,
1607+
CIRIfOpLowering, CIRVectorCreateLowering, CIRVectorInsertLowering,
1608+
CIRVectorExtractLowering, CIRVectorCmpOpLowering, CIRACosOpLowering,
1609+
CIRASinOpLowering, CIRUnreachableOpLowering, CIRTanOpLowering,
1610+
CIRTrapOpLowering>(converter, patterns.getContext());
15221611
}
15231612

15241613
static mlir::TypeConverter prepareTypeConverter() {
@@ -1624,6 +1713,7 @@ mlir::ModuleOp lowerFromCIRToMLIRToLLVMDialect(mlir::ModuleOp theModule,
16241713

16251714
mlir::PassManager pm(mlirCtx);
16261715

1716+
pm.addPass(createMLIRCoreDialectsLoweringPreparePass());
16271717
pm.addPass(createConvertCIRToMLIRPass());
16281718
pm.addPass(createConvertMLIRToLLVMPass());
16291719

@@ -1669,6 +1759,7 @@ mlir::ModuleOp lowerFromCIRToMLIR(mlir::ModuleOp theModule,
16691759

16701760
mlir::PassManager pm(mlirCtx);
16711761

1762+
pm.addPass(createMLIRCoreDialectsLoweringPreparePass());
16721763
pm.addPass(createConvertCIRToMLIRPass());
16731764

16741765
auto result = !mlir::failed(pm.run(theModule));

‎clang/lib/CIR/Lowering/ThroughMLIR/MLIRCoreDialectsLoweringPrepare.cpp

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
//===- MLIRCoreDialectsLoweringPrepare.cpp - CIR lowering preparation -----===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#include "mlir/IR/BuiltinOps.h"
10+
#include "mlir/IR/IRMapping.h"
11+
#include "mlir/Pass/Pass.h"
12+
#include "mlir/Transforms/DialectConversion.h"
13+
#include "clang/CIR/Dialect/Builder/CIRBaseBuilder.h"
14+
#include "clang/CIR/Dialect/IR/CIRDialect.h"
15+
16+
using namespace llvm;
17+
using namespace cir;
18+
19+
namespace cir {
20+
21+
struct MLIRLoweringPrepare
22+
: public mlir::PassWrapper<MLIRLoweringPrepare,
23+
mlir::OperationPass<mlir::ModuleOp>> {
24+
// `scf.index_switch` requires that switch branches do not fall through.
25+
// We need to copy the next branch's body when the current `cir.case` does
26+
// not terminate with a break.
27+
void removeFallthrough(llvm::SmallVector<CaseOp> &cases);
28+
29+
void runOnOp(mlir::Operation *op);
30+
void runOnOperation() final;
31+
32+
StringRef getDescription() const override {
33+
return "Rewrite CIR module to be more 'scf' dialect-friendly";
34+
}
35+
36+
StringRef getArgument() const override { return "mlir-lowering-prepare"; }
37+
};
38+
39+
// `scf.index_switch` requires that switch branches do not fall through.
40+
// We need to copy the next branch's body when the current `cir.case` does not
41+
// terminate with a break.
42+
void MLIRLoweringPrepare::removeFallthrough(llvm::SmallVector<CaseOp> &cases) {
43+
CIRBaseBuilderTy builder(getContext());
44+
// Note we enumerate in the reverse order, to facilitate the cloning.
45+
for (auto it = cases.rbegin(); it != cases.rend(); it++) {
46+
auto caseOp = *it;
47+
auto &region = caseOp.getRegion();
48+
auto &lastBlock = region.back();
49+
mlir::Operation &last = lastBlock.back();
50+
if (isa<BreakOp>(last))
51+
continue;
52+
53+
// The last op must be a `cir.yield`. As it falls through, we copy the
54+
// previous case's body to this one.
55+
if (!isa<YieldOp>(last)) {
56+
caseOp->dump();
57+
continue;
58+
}
59+
assert(isa<YieldOp>(last));
60+
61+
// If there's no previous case, we can simply change the yield into a break.
62+
if (it == cases.rbegin()) {
63+
builder.setInsertionPointAfter(&last);
64+
builder.create<BreakOp>(last.getLoc());
65+
last.erase();
66+
continue;
67+
}
68+
69+
auto prevIt = it;
70+
--prevIt;
71+
CaseOp &prev = *prevIt;
72+
auto &prevRegion = prev.getRegion();
73+
mlir::IRMapping mapping;
74+
builder.cloneRegionBefore(prevRegion, region, region.end());
75+
76+
// We inline the block to the end.
77+
// This is required because `scf.index_switch` expects that each of its
78+
// region contains a single block.
79+
mlir::Block *cloned = lastBlock.getNextNode();
80+
for (auto it = cloned->begin(); it != cloned->end();) {
81+
auto next = it;
82+
next++;
83+
it->moveBefore(&last);
84+
it = next;
85+
}
86+
cloned->erase();
87+
last.erase();
88+
}
89+
}
90+
91+
void MLIRLoweringPrepare::runOnOp(mlir::Operation *op) {
92+
if (auto switchOp = dyn_cast<SwitchOp>(op)) {
93+
llvm::SmallVector<CaseOp> cases;
94+
if (!switchOp.isSimpleForm(cases))
95+
op->emitError("NYI");
96+
97+
removeFallthrough(cases);
98+
return;
99+
}
100+
op->emitError("unexpected op type");
101+
}
102+
103+
void MLIRLoweringPrepare::runOnOperation() {
104+
auto module = getOperation();
105+
106+
llvm::SmallVector<mlir::Operation *> opsToTransform;
107+
module->walk([&](mlir::Operation *op) {
108+
if (isa<SwitchOp>(op))
109+
opsToTransform.push_back(op);
110+
});
111+
112+
for (auto *op : opsToTransform)
113+
runOnOp(op);
114+
}
115+
116+
std::unique_ptr<mlir::Pass> createMLIRCoreDialectsLoweringPreparePass() {
117+
return std::make_unique<MLIRLoweringPrepare>();
118+
}
119+
120+
} // namespace cir

‎clang/test/CIR/Lowering/ThroughMLIR/switch.c

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -fclangir -fno-clangir-direct-lowering -emit-mlir=core %s -o %t.mlir
2+
// RUN: FileCheck --input-file=%t.mlir %s
3+
4+
void fallthrough() {
5+
int i = 0;
6+
switch (i) {
7+
case 2:
8+
i++;
9+
case 3:
10+
i++;
11+
break;
12+
case 8:
13+
i++;
14+
}
15+
16+
// This should copy the `i++; break` in case 3 to case 2.
17+
18+
// CHECK: memref.alloca_scope {
19+
// CHECK: %[[I:.+]] = memref.load %alloca[]
20+
// CHECK: %[[CASTED:.+]] = arith.index_cast %[[I]]
21+
// CHECK: scf.index_switch %[[CASTED]]
22+
// CHECK: case 2 {
23+
// CHECK: %[[I:.+]] = memref.load %alloca[]
24+
// CHECK: %[[ONE:.+]] = arith.constant 1
25+
// CHECK: %[[ADD:.+]] = arith.addi %[[I]], %[[ONE]]
26+
// CHECK: memref.store %[[ADD]], %alloca[]
27+
// CHECK: %[[I:.+]] = memref.load %alloca[]
28+
// CHECK: %[[ONE:.+]] = arith.constant 1
29+
// CHECK: %[[ADD:.+]] = arith.addi %[[I]], %[[ONE]]
30+
// CHECK: memref.store %[[ADD]], %alloca[]
31+
// CHECK: scf.yield
32+
// CHECK: }
33+
// CHECK: case 3 {
34+
// CHECK: %[[I:.+]] = memref.load %alloca[]
35+
// CHECK: %[[ONE:.+]] = arith.constant 1
36+
// CHECK: %[[ADD:.+]] = arith.addi %[[I]], %[[ONE]]
37+
// CHECK: memref.store %[[ADD]], %alloca[]
38+
// CHECK: scf.yield
39+
// CHECK: }
40+
// CHECK: case 8 {
41+
// CHECK: %[[I:.+]] = memref.load %alloca[]
42+
// CHECK: %[[ONE:.+]] = arith.constant 1
43+
// CHECK: %[[ADD:.+]] = arith.addi %[[I]], %[[ONE]]
44+
// CHECK: memref.store %[[ADD]], %alloca[]
45+
// CHECK: scf.yield
46+
// CHECK: }
47+
// CHECK: default {
48+
// CHECK: }
49+
// CHECK: }
50+
}

0 commit comments

Comments
 (0)