[Mlir-commits] [mlir] [mlir] [tblgen-to-irdl] Add region support (PR #110512)

[Alex Rice]

https://github.com/alexarice created https://github.com/llvm/llvm-project/pull/110512

Adds support to exporting regions to irdl

@math-fehr 

>From 3e573c8ec5d32ee06d584d79cf29e7daeea1eec3 Mon Sep 17 00:00:00 2001
From: Alex Rice <alexrice999 at hotmail.co.uk>
Date: Thu, 15 Aug 2024 15:17:03 +0100
Subject: [PATCH] Add region support

---
 mlir/include/mlir/IR/OpBase.td                |  4 ++-
 mlir/test/tblgen-to-irdl/TestDialect.td       | 11 +++++++
 .../tools/tblgen-to-irdl/OpDefinitionsGen.cpp | 31 +++++++++++++++++++
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 4481e56615b8bf..5c82c041c62eeb 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -206,7 +206,9 @@ def AnyRegion : Region<CPred<"true">, "any region">;
 // A region with the given number of blocks.
 class SizedRegion<int numBlocks> : Region<
   CPred<"::llvm::hasNItems($_self, " # numBlocks # ")">,
-  "region with " # numBlocks # " blocks">;
+  "region with " # numBlocks # " blocks"> {
+  int blocks = numBlocks;
+}
 
 // A region with at least the given number of blocks.
 class MinSizedRegion<int numBlocks> : Region<
diff --git a/mlir/test/tblgen-to-irdl/TestDialect.td b/mlir/test/tblgen-to-irdl/TestDialect.td
index 1ba84a5d3683d4..7f4815d865b60b 100644
--- a/mlir/test/tblgen-to-irdl/TestDialect.td
+++ b/mlir/test/tblgen-to-irdl/TestDialect.td
@@ -106,6 +106,17 @@ def Test_OrOp : Test_Op<"or"> {
 // CHECK-NEXT:    irdl.operands(%[[v3]])
 // CHECK-NEXT:  }
 
+// Check regions are converted correctly.
+def Test_RegionsOp : Test_Op<"regions"> {
+  let regions = (region AnyRegion:$any_region,
+                        SizedRegion<1>:$single_block_region);
+}
+// CHECK-LABEL: irdl.operation @regions {
+// CHECK-NEXT:    %[[v0:[^ ]*]] = irdl.region
+// CHECK-NEXT:    %[[v1:[^ ]*]] = irdl.region with size 1
+// CHECK-NEXT:    irdl.regions(%[[v0]], %[[v1]])
+// CHECK-NEXT:  }
+
 // Check that various types are converted correctly.
 def Test_TypesOp : Test_Op<"types"> {
   let arguments = (ins I32:$a,
diff --git a/mlir/tools/tblgen-to-irdl/OpDefinitionsGen.cpp b/mlir/tools/tblgen-to-irdl/OpDefinitionsGen.cpp
index d0a3552fb123da..79ff919f634b02 100644
--- a/mlir/tools/tblgen-to-irdl/OpDefinitionsGen.cpp
+++ b/mlir/tools/tblgen-to-irdl/OpDefinitionsGen.cpp
@@ -338,6 +338,29 @@ Value createAttrConstraint(OpBuilder &builder, tblgen::Constraint constraint) {
   return createPredicate(builder, constraint.getPredicate());
 }
 
+Value createRegionConstraint(OpBuilder &builder, tblgen::Region constraint) {
+  MLIRContext *ctx = builder.getContext();
+  const Record &predRec = constraint.getDef();
+
+  if (predRec.getName() == "AnyRegion") {
+    ValueRange entryBlockArgs = {};
+    auto op =
+        builder.create<irdl::RegionOp>(UnknownLoc::get(ctx), entryBlockArgs);
+    return op.getResult();
+  }
+
+  if (predRec.isSubClassOf("SizedRegion")) {
+    ValueRange entryBlockArgs = {};
+    auto ty = IntegerType::get(ctx, 32);
+    auto op = builder.create<irdl::RegionOp>(
+        UnknownLoc::get(ctx), entryBlockArgs,
+        IntegerAttr::get(ty, predRec.getValueAsInt("blocks")));
+    return op.getResult();
+  }
+
+  return createPredicate(builder, constraint.getPredicate());
+}
+
 /// Returns the name of the operation without the dialect prefix.
 static StringRef getOperatorName(tblgen::Operator &tblgenOp) {
   StringRef opName = tblgenOp.getDef().getValueAsString("opName");
@@ -404,6 +427,12 @@ irdl::OperationOp createIRDLOperation(OpBuilder &builder,
     attrNames.push_back(StringAttr::get(ctx, namedAttr.name));
   }
 
+  SmallVector<Value> regions;
+  for (auto namedRegion : tblgenOp.getRegions()) {
+    regions.push_back(
+        createRegionConstraint(consBuilder, namedRegion.constraint));
+  }
+
   // Create the operands and results operations.
   if (!operands.empty())
     consBuilder.create<irdl::OperandsOp>(UnknownLoc::get(ctx), operands,
@@ -414,6 +443,8 @@ irdl::OperationOp createIRDLOperation(OpBuilder &builder,
   if (!attributes.empty())
     consBuilder.create<irdl::AttributesOp>(UnknownLoc::get(ctx), attributes,
                                            ArrayAttr::get(ctx, attrNames));
+  if (!regions.empty())
+    consBuilder.create<irdl::RegionsOp>(UnknownLoc::get(ctx), regions);
 
   return op;
 }