[CIR][ThroughMLIR] Lower simple SwitchOp

clang/include/clang/CIR/Passes.h

@@ -18,6 +18,11 @@
 #include <memory>
 
 namespace cir {
+/// Create a pass for transforming CIR operations to more 'scf' dialect-friendly
+/// forms. It rewrites operations that aren't supported by 'scf', such as breaks
+/// and continues.
+std::unique_ptr<mlir::Pass> createMLIRCoreDialectsLoweringPreparePass();
+
 /// Create a pass for lowering from MLIR builtin dialects such as `Affine` and
 /// `Std`, to the LLVM dialect for codegen.
 std::unique_ptr<mlir::Pass> createConvertMLIRToLLVMPass();

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

@@ -9,6 +9,7 @@ add_clang_library(clangCIRLoweringThroughMLIR
   LowerCIRLoopToSCF.cpp
   LowerCIRToMLIR.cpp
   LowerMLIRToLLVM.cpp
+  MLIRCoreDialectsLoweringPrepare.cpp
 
   DEPENDS
   MLIRCIROpsIncGen

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

@@ -1496,29 +1496,118 @@ class CIRTrapOpLowering : public mlir::OpConversionPattern<cir::TrapOp> {
   }
 };
 
+class CIRSwitchOpLowering : public mlir::OpConversionPattern<cir::SwitchOp> {
+public:
+  using OpConversionPattern<cir::SwitchOp>::OpConversionPattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(cir::SwitchOp op, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    rewriter.setInsertionPointAfter(op);
+    llvm::SmallVector<CaseOp> cases;
+    if (!op.isSimpleForm(cases))
+      llvm_unreachable("NYI");
+
+    llvm::SmallVector<int64_t> caseValues;
+    // Maps the index of a CaseOp in `cases`, to the index in `caseValues`.
+    // This is necessary because some CaseOp might carry 0 or multiple values.
+    llvm::DenseMap<size_t, unsigned> indexMap;
+    caseValues.reserve(cases.size());
+    for (auto [i, caseOp] : llvm::enumerate(cases)) {
+      switch (caseOp.getKind()) {
+      case CaseOpKind::Equal: {
+        auto valueAttr = caseOp.getValue()[0];
+        auto value = cast<cir::IntAttr>(valueAttr);
+        indexMap[i] = caseValues.size();
+        caseValues.push_back(value.getUInt());
+        break;
+      }
+      case CaseOpKind::Default:
+        break;
+      case CaseOpKind::Range:
+      case CaseOpKind::Anyof:
+        llvm_unreachable("NYI");
+      }
+    }
+
+    auto operand = adaptor.getOperands()[0];
+    // `scf.index_switch` expects an index of type `index`.
+    auto indexType = mlir::IndexType::get(getContext());
+    auto indexCast = rewriter.create<mlir::arith::IndexCastOp>(
+        op.getLoc(), indexType, operand);
+    auto indexSwitch = rewriter.create<mlir::scf::IndexSwitchOp>(
+        op.getLoc(), mlir::TypeRange{}, indexCast, caseValues, cases.size());
+
+    bool metDefault = false;
+    for (auto [i, caseOp] : llvm::enumerate(cases)) {
+      auto &region = caseOp.getRegion();
+      switch (caseOp.getKind()) {
+      case CaseOpKind::Equal: {
+        auto &caseRegion = indexSwitch.getCaseRegions()[indexMap[i]];
+        rewriter.inlineRegionBefore(region, caseRegion, caseRegion.end());
+        break;
+      }
+      case CaseOpKind::Default: {
+        auto &defaultRegion = indexSwitch.getDefaultRegion();
+        rewriter.inlineRegionBefore(region, defaultRegion, defaultRegion.end());
+        metDefault = true;
+        break;
+      }
+      case CaseOpKind::Range:
+      case CaseOpKind::Anyof:
+        llvm_unreachable("NYI");
+      }
+    }
+
+    // `scf.index_switch` expects its default region to contain exactly one
+    // block. If we don't have a default region in `cir.switch`, we need to
+    // supply it here.
+    if (!metDefault) {
+      auto &defaultRegion = indexSwitch.getDefaultRegion();
+      mlir::Block *block =
+          rewriter.createBlock(&defaultRegion, defaultRegion.end());
+      rewriter.setInsertionPointToEnd(block);
+      rewriter.create<mlir::scf::YieldOp>(op.getLoc());
+    }
+
+    // The final `cir.break` should be replaced to `scf.yield`.
+    // After MLIRLoweringPrepare pass, every case must end with a `cir.break`.
+    for (auto &region : indexSwitch.getCaseRegions()) {
+      auto &lastBlock = region.back();
+      auto &lastOp = lastBlock.back();
+      assert(isa<BreakOp>(lastOp));
+      rewriter.setInsertionPointAfter(&lastOp);
+      rewriter.replaceOpWithNewOp<mlir::scf::YieldOp>(&lastOp);
+    }
+
+    rewriter.replaceOp(op, indexSwitch);
+
+    return mlir::success();
+  }
+};
+
 void populateCIRToMLIRConversionPatterns(mlir::RewritePatternSet &patterns,
                                          mlir::TypeConverter &converter) {
   patterns.add<CIRReturnLowering, CIRBrOpLowering>(patterns.getContext());
 
-  patterns
-      .add<CIRATanOpLowering, CIRCmpOpLowering, CIRCallOpLowering,
-           CIRUnaryOpLowering, CIRBinOpLowering, CIRLoadOpLowering,
-           CIRConstantOpLowering, CIRStoreOpLowering, CIRAllocaOpLowering,
-           CIRFuncOpLowering, CIRScopeOpLowering, CIRBrCondOpLowering,
-           CIRTernaryOpLowering, CIRYieldOpLowering, CIRCosOpLowering,
-           CIRGlobalOpLowering, CIRGetGlobalOpLowering, CIRCastOpLowering,
-           CIRPtrStrideOpLowering, CIRGetElementOpLowering, CIRSqrtOpLowering,
-           CIRCeilOpLowering, CIRExp2OpLowering, CIRExpOpLowering,
-           CIRFAbsOpLowering, CIRAbsOpLowering, CIRFloorOpLowering,
-           CIRLog10OpLowering, CIRLog2OpLowering, CIRLogOpLowering,
-           CIRRoundOpLowering, CIRPtrStrideOpLowering, CIRSinOpLowering,
-           CIRShiftOpLowering, CIRBitClzOpLowering, CIRBitCtzOpLowering,
-           CIRBitPopcountOpLowering, CIRBitClrsbOpLowering, CIRBitFfsOpLowering,
-           CIRBitParityOpLowering, CIRIfOpLowering, CIRVectorCreateLowering,
-           CIRVectorInsertLowering, CIRVectorExtractLowering,
-           CIRVectorCmpOpLowering, CIRACosOpLowering, CIRASinOpLowering,
-           CIRUnreachableOpLowering, CIRTanOpLowering, CIRTrapOpLowering>(
-          converter, patterns.getContext());
+  patterns.add<
+      CIRSwitchOpLowering, CIRGetElementOpLowering, CIRATanOpLowering,
+      CIRCmpOpLowering, CIRCallOpLowering, CIRUnaryOpLowering, CIRBinOpLowering,
+      CIRLoadOpLowering, CIRConstantOpLowering, CIRStoreOpLowering,
+      CIRAllocaOpLowering, CIRFuncOpLowering, CIRScopeOpLowering,
+      CIRBrCondOpLowering, CIRTernaryOpLowering, CIRYieldOpLowering,
+      CIRCosOpLowering, CIRGlobalOpLowering, CIRGetGlobalOpLowering,
+      CIRCastOpLowering, CIRPtrStrideOpLowering, CIRSqrtOpLowering,
+      CIRCeilOpLowering, CIRExp2OpLowering, CIRExpOpLowering, CIRFAbsOpLowering,
+      CIRAbsOpLowering, CIRFloorOpLowering, CIRLog10OpLowering,
+      CIRLog2OpLowering, CIRLogOpLowering, CIRRoundOpLowering,
+      CIRPtrStrideOpLowering, CIRSinOpLowering, CIRShiftOpLowering,
+      CIRBitClzOpLowering, CIRBitCtzOpLowering, CIRBitPopcountOpLowering,
+      CIRBitClrsbOpLowering, CIRBitFfsOpLowering, CIRBitParityOpLowering,
+      CIRIfOpLowering, CIRVectorCreateLowering, CIRVectorInsertLowering,
+      CIRVectorExtractLowering, CIRVectorCmpOpLowering, CIRACosOpLowering,
+      CIRASinOpLowering, CIRUnreachableOpLowering, CIRTanOpLowering,
+      CIRTrapOpLowering>(converter, patterns.getContext());
 }
 
 static mlir::TypeConverter prepareTypeConverter() {
@@ -1624,6 +1713,7 @@ mlir::ModuleOp lowerFromCIRToMLIRToLLVMDialect(mlir::ModuleOp theModule,
 
   mlir::PassManager pm(mlirCtx);
 
+  pm.addPass(createMLIRCoreDialectsLoweringPreparePass());
   pm.addPass(createConvertCIRToMLIRPass());
   pm.addPass(createConvertMLIRToLLVMPass());
 
@@ -1669,6 +1759,7 @@ mlir::ModuleOp lowerFromCIRToMLIR(mlir::ModuleOp theModule,
 
   mlir::PassManager pm(mlirCtx);
 
+  pm.addPass(createMLIRCoreDialectsLoweringPreparePass());
   pm.addPass(createConvertCIRToMLIRPass());
 
   auto result = !mlir::failed(pm.run(theModule));

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

@@ -0,0 +1,120 @@
+//===- MLIRCoreDialectsLoweringPrepare.cpp - CIR lowering preparation -----===//
+//
+// 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 "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "clang/CIR/Dialect/Builder/CIRBaseBuilder.h"
+#include "clang/CIR/Dialect/IR/CIRDialect.h"
+
+using namespace llvm;
+using namespace cir;
+
+namespace cir {
+
+struct MLIRLoweringPrepare
+    : public mlir::PassWrapper<MLIRLoweringPrepare,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+  // `scf.index_switch` requires that switch branches do not fall through.
+  // We need to copy the next branch's body when the current `cir.case` does
+  // not terminate with a break.
+  void removeFallthrough(llvm::SmallVector<CaseOp> &cases);
+
+  void runOnOp(mlir::Operation *op);
+  void runOnOperation() final;
+
+  StringRef getDescription() const override {
+    return "Rewrite CIR module to be more 'scf' dialect-friendly";
+  }
+
+  StringRef getArgument() const override { return "mlir-lowering-prepare"; }
+};
+
+// `scf.index_switch` requires that switch branches do not fall through.
+// We need to copy the next branch's body when the current `cir.case` does not
+// terminate with a break.
+void MLIRLoweringPrepare::removeFallthrough(llvm::SmallVector<CaseOp> &cases) {
+  CIRBaseBuilderTy builder(getContext());
+  // Note we enumerate in the reverse order, to facilitate the cloning.
+  for (auto it = cases.rbegin(); it != cases.rend(); it++) {
+    auto caseOp = *it;
+    auto &region = caseOp.getRegion();
+    auto &lastBlock = region.back();
+    mlir::Operation &last = lastBlock.back();
+    if (isa<BreakOp>(last))
+      continue;
+
+    // The last op must be a `cir.yield`. As it falls through, we copy the
+    // previous case's body to this one.
+    if (!isa<YieldOp>(last)) {
+      caseOp->dump();
+      continue;
+    }
+    assert(isa<YieldOp>(last));
+
+    // If there's no previous case, we can simply change the yield into a break.
+    if (it == cases.rbegin()) {
+      builder.setInsertionPointAfter(&last);
+      builder.create<BreakOp>(last.getLoc());
+      last.erase();
+      continue;
+    }
+
+    auto prevIt = it;
+    --prevIt;
+    CaseOp &prev = *prevIt;
+    auto &prevRegion = prev.getRegion();
+    mlir::IRMapping mapping;
+    builder.cloneRegionBefore(prevRegion, region, region.end());
+
+    // We inline the block to the end.
+    // This is required because `scf.index_switch` expects that each of its
+    // region contains a single block.
+    mlir::Block *cloned = lastBlock.getNextNode();
+    for (auto it = cloned->begin(); it != cloned->end();) {
+      auto next = it;
+      next++;
+      it->moveBefore(&last);
+      it = next;
+    }
+    cloned->erase();
+    last.erase();
+  }
+}
+
+void MLIRLoweringPrepare::runOnOp(mlir::Operation *op) {
+  if (auto switchOp = dyn_cast<SwitchOp>(op)) {
+    llvm::SmallVector<CaseOp> cases;
+    if (!switchOp.isSimpleForm(cases))
+      op->emitError("NYI");
+
+    removeFallthrough(cases);
+    return;
+  }
+  op->emitError("unexpected op type");
+}
+
+void MLIRLoweringPrepare::runOnOperation() {
+  auto module = getOperation();
+
+  llvm::SmallVector<mlir::Operation *> opsToTransform;
+  module->walk([&](mlir::Operation *op) {
+    if (isa<SwitchOp>(op))
+      opsToTransform.push_back(op);
+  });
+
+  for (auto *op : opsToTransform)
+    runOnOp(op);
+}
+
+std::unique_ptr<mlir::Pass> createMLIRCoreDialectsLoweringPreparePass() {
+  return std::make_unique<MLIRLoweringPrepare>();
+}
+
+} // namespace cir

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

@@ -0,0 +1,50 @@
+// RUN: %clang_cc1 -triple x86_64-unknown-linux-gnu -fclangir -fno-clangir-direct-lowering -emit-mlir=core %s -o %t.mlir
+// RUN: FileCheck --input-file=%t.mlir %s
+
+void fallthrough() {
+  int i = 0;
+  switch (i) {
+  case 2:
+    i++;
+  case 3:
+    i++;
+    break;
+  case 8:
+    i++;
+  }
+
+  // This should copy the `i++; break` in case 3 to case 2.
+
+  // CHECK: memref.alloca_scope  {
+  // CHECK:   %[[I:.+]] = memref.load %alloca[]
+  // CHECK:   %[[CASTED:.+]] = arith.index_cast %[[I]]
+  // CHECK:   scf.index_switch %[[CASTED]]
+  // CHECK:   case 2 {
+  // CHECK:     %[[I:.+]] = memref.load %alloca[]
+  // CHECK:     %[[ONE:.+]] = arith.constant 1
+  // CHECK:     %[[ADD:.+]] = arith.addi %[[I]], %[[ONE]]
+  // CHECK:     memref.store %[[ADD]], %alloca[]
+  // CHECK:     %[[I:.+]] = memref.load %alloca[]
+  // CHECK:     %[[ONE:.+]] = arith.constant 1
+  // CHECK:     %[[ADD:.+]] = arith.addi %[[I]], %[[ONE]]
+  // CHECK:     memref.store %[[ADD]], %alloca[]
+  // CHECK:     scf.yield
+  // CHECK:   }
+  // CHECK:   case 3 {
+  // CHECK:     %[[I:.+]] = memref.load %alloca[]
+  // CHECK:     %[[ONE:.+]] = arith.constant 1
+  // CHECK:     %[[ADD:.+]] = arith.addi %[[I]], %[[ONE]]
+  // CHECK:     memref.store %[[ADD]], %alloca[]
+  // CHECK:     scf.yield
+  // CHECK:   }
+  // CHECK:   case 8 {
+  // CHECK:     %[[I:.+]] = memref.load %alloca[]
+  // CHECK:     %[[ONE:.+]] = arith.constant 1
+  // CHECK:     %[[ADD:.+]] = arith.addi %[[I]], %[[ONE]]
+  // CHECK:     memref.store %[[ADD]], %alloca[]
+  // CHECK:     scf.yield
+  // CHECK:   }
+  // CHECK:   default {
+  // CHECK:   }
+  // CHECK: }
+}