[Mlir-commits] [mlir] [mlir] [scf] Add index_switch support for SCFStructuralTypeConversion (PR #118588)

[llvmlistbot at llvm.org]

https://github.com/Menooker created https://github.com/llvm/llvm-project/pull/118588

Add the support of index_switch which is missing in SCFStructuralTypeConversion. 

>From ed4e517f39e90873abc63022f1ed142be272468c Mon Sep 17 00:00:00 2001
From: Menooker <myjisgreat at live.cn>
Date: Tue, 3 Dec 2024 19:58:24 -0800
Subject: [PATCH] [mlir] [scf] Add index_switch support for
 SCFStructuralTypeConversion

---
 .../Transforms/StructuralTypeConversions.cpp  | 34 ++++++++++++++++---
 1 file changed, 30 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp b/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp
index c0589044c26ecb..3edcdbcf84bd05 100644
--- a/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp
@@ -190,6 +190,32 @@ class ConvertWhileOpTypes
 };
 } // namespace
 
+namespace {
+class ConvertIndexSwitchOpTypes
+    : public Structural1ToNConversionPattern<IndexSwitchOp,
+                                             ConvertIndexSwitchOpTypes> {
+public:
+  using Structural1ToNConversionPattern::Structural1ToNConversionPattern;
+
+  std::optional<IndexSwitchOp>
+  convertSourceOp(IndexSwitchOp op, OneToNOpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter,
+                  TypeRange dstTypes) const {
+    auto newOp = rewriter.create<scf::IndexSwitchOp>(
+        op.getLoc(), dstTypes,
+        getSingleValue(flattenValues(adaptor.getOperands())),
+        adaptor.getCases(),
+        /*numCaseRegions=*/adaptor.getCaseRegions().size());
+    newOp->setAttrs(op->getAttrs());
+    for (auto &&[newRegion, oldRegion] :
+         llvm::zip_equal(newOp.getRegions(), adaptor.getRegions())) {
+      rewriter.inlineRegionBefore(*oldRegion, *newRegion, newRegion->end());
+    }
+    return newOp;
+  }
+};
+} // namespace
+
 namespace {
 // When the result types of a ForOp/IfOp get changed, the operand types of the
 // corresponding yield op need to be changed. In order to trigger the
@@ -224,19 +250,19 @@ class ConvertConditionOpTypes : public OpConversionPattern<ConditionOp> {
 void mlir::scf::populateSCFStructuralTypeConversions(
     const TypeConverter &typeConverter, RewritePatternSet &patterns) {
   patterns.add<ConvertForOpTypes, ConvertIfOpTypes, ConvertYieldOpTypes,
-               ConvertWhileOpTypes, ConvertConditionOpTypes>(
-      typeConverter, patterns.getContext());
+               ConvertWhileOpTypes, ConvertIndexSwitchOpTypes,
+               ConvertConditionOpTypes>(typeConverter, patterns.getContext());
 }
 
 void mlir::scf::populateSCFStructuralTypeConversionTarget(
     const TypeConverter &typeConverter, ConversionTarget &target) {
-  target.addDynamicallyLegalOp<ForOp, IfOp>([&](Operation *op) {
+  target.addDynamicallyLegalOp<ForOp, IfOp, IndexSwitchOp>([&](Operation *op) {
     return typeConverter.isLegal(op->getResultTypes());
   });
   target.addDynamicallyLegalOp<scf::YieldOp>([&](scf::YieldOp op) {
     // We only have conversions for a subset of ops that use scf.yield
     // terminators.
-    if (!isa<ForOp, IfOp, WhileOp>(op->getParentOp()))
+    if (!isa<ForOp, IfOp, WhileOp, IndexSwitchOp>(op->getParentOp()))
       return true;
     return typeConverter.isLegal(op.getOperandTypes());
   });