[mlir][mlir-pdll][Parser] Allow parsing dictionaries in Constraint se…

…ction
Xilinx · Apr 9, 2024 · 3280a80 · 3280a80
1 parent df673ea
commit 3280a80
Show file tree

Hide file tree

Showing 3 changed files with 101 additions and 10 deletions.
diff --git a/mlir/lib/Tools/PDLL/Parser/Parser.cpp b/mlir/lib/Tools/PDLL/Parser/Parser.cpp
@@ -600,8 +600,10 @@ class Parser {
   CodeCompleteContext *codeCompleteContext;
 
   struct {
-    ast::UserRewriteDecl *createDictionaryAttr;
-    ast::UserRewriteDecl *addEntryToDictionaryAttr;
+    ast::UserRewriteDecl *createDictionaryAttr_Rewrite;
+    ast::UserConstraintDecl *createDictionaryAttr_Constraint;
+    ast::UserRewriteDecl *addEntryToDictionaryAttr_Rewrite;
+    ast::UserConstraintDecl *addEntryToDictionaryAttr_Constraint;
     ast::UserRewriteDecl *createArrayAttr;
     ast::UserRewriteDecl *addElemToArrayAttr;
     ast::UserConstraintDecl *add;
@@ -655,7 +657,7 @@ void Parser::declareBuiltins() {
       declareBuiltin<ast::UserConstraintDecl>(
           "__builtin_addEntryToDictionaryAttr_constraint",
           {"attr", "attrName", "attrEntry"},
-      /*returnsAttr=*/true);
+          /*returnsAttr=*/true);
   builtins.createArrayAttr = declareBuiltin<ast::UserRewriteDecl>(
       "__builtin_createArrayAttr", {}, /*returnsAttr=*/true);
   builtins.addElemToArrayAttr = declareBuiltin<ast::UserRewriteDecl>(
@@ -2121,14 +2123,22 @@ FailureOr<ast::Expr *> Parser::parseDictAttrExpr() {
   consumeToken(Token::l_brace);
   SMRange loc = curToken.getLoc();
 
-  if (parserContext != ParserContext::Rewrite)
-    return emitError(
-        "Parsing of dictionary attributes as constraint not supported!");
-
-  auto dictAttrCall = createBuiltinCall(loc, builtins.createDictionaryAttr, {});
+  FailureOr<ast::Expr *> dictAttrCall;
+  if (parserContext == ParserContext::Rewrite) {
+    dictAttrCall =
+        createBuiltinCall(loc, builtins.createDictionaryAttr_Rewrite, {});
+  } else {
+    dictAttrCall =
+        createBuiltinCall(loc, builtins.createDictionaryAttr_Constraint, {});
+  }
   if (failed(dictAttrCall))
     return failure();
 
+  // No key-values inside dictionary
+  if (consumeIf(Token::r_brace)) {
+    return dictAttrCall;
+  }
+
   // Add each nested attribute to the dict
   do {
     FailureOr<ast::NamedAttributeDecl *> decl =
@@ -2159,8 +2169,15 @@ FailureOr<ast::Expr *> Parser::parseDictAttrExpr() {
     // Create addEntryToDictionaryAttr native call.
     SmallVector<ast::Expr *> arrayAttrArgs{*dictAttrCall, *stringAttrRef,
                                            namedDecl->getValue()};
-    auto entryToDictionaryCall = createBuiltinCall(
-        loc, builtins.addEntryToDictionaryAttr, arrayAttrArgs);
+
+    FailureOr<ast::Expr *> entryToDictionaryCall;
+    if (parserContext == ParserContext::Rewrite) {
+      entryToDictionaryCall = createBuiltinCall(
+          loc, builtins.addEntryToDictionaryAttr_Rewrite, arrayAttrArgs);
+    } else {
+      entryToDictionaryCall = createBuiltinCall(
+          loc, builtins.addEntryToDictionaryAttr_Constraint, arrayAttrArgs);
+    }
     if (failed(entryToDictionaryCall))
       return failure();
 

diff --git a/mlir/test/mlir-pdll/CodeGen/MLIR/expr.pdll b/mlir/test/mlir-pdll/CodeGen/MLIR/expr.pdll
@@ -253,6 +253,36 @@ Pattern RewriteMultiplyElementsArrayAttr {
   };
 }
 
+// -----
+
+// CHECK-LABEL:    pdl.pattern @ConstraintWithEmptyDictionary : benefit(0) {
+//        CHECK:    %[[DICT:.+]] = apply_native_constraint "__builtin_createDictionaryAttr_constraint" : !pdl.attribute
+//        CHECK:    rewrite %{{.*}} {
+//   CHECK-NEXT:      operation "test.success"  {"importantAttr" = %[[DICT]]}
+
+Pattern ConstraintWithEmptyDictionary {
+  let dict = {};
+  replace op<test.op> with op<test.success>() { importantAttr = dict };
+}
+
+// -----
+// CHECK-LABEL:    pdl.pattern @ConstraintWithTwoEntriesDictionary : benefit(0) {
+//        CHECK:    %[[DICT:.+]] = apply_native_constraint "__builtin_createDictionaryAttr_constraint" : !pdl.attribute
+//   CHECK-NEXT:    %[[KEY:.+]] = attribute = "hello"
+//   CHECK-NEXT:    %[[VAL:.+]] = attribute = "world"
+//   CHECK-NEXT:    %[[DICT2:.+]] = apply_native_constraint "__builtin_addEntryToDictionaryAttr_constraint"(%[[DICT]], %[[KEY]], %[[VAL]] : !pdl.attribute, !pdl.attribute, !pdl.attribute) : !pdl.attribute
+//   CHECK-NEXT:    %[[KEY2:.+]] = attribute = "bye"
+//   CHECK-NEXT:    %[[VAL2:.+]] = attribute = 100 : ui8
+//   CHECK-NEXT:    %[[DICT3:.+]] = apply_native_constraint "__builtin_addEntryToDictionaryAttr_constraint"(%[[DICT2]], %[[KEY2]], %[[VAL2]] : !pdl.attribute, !pdl.attribute, !pdl.attribute) : !pdl.attribute
+//        CHECK:    rewrite %{{.*}} {
+//   CHECK-NEXT:      operation "test.success"  {"importantAttr" = %[[DICT3]]}
+
+
+Pattern ConstraintWithTwoEntriesDictionary {
+  let dict = { hello = "world", bye = 100 : ui8 };
+  replace op<test.op> with op<test.success>() { importantAttr = dict };
+}
+
 // -----
 // CHECK-LABEL:    pdl.pattern @TestAdd : benefit(0) {
 // CHECK:            %[[VAL_0:.*]] = attribute = 4 : i32

diff --git a/mlir/test/mlir-pdll/Parser/expr.pdll b/mlir/test/mlir-pdll/Parser/expr.pdll
@@ -55,6 +55,50 @@ Pattern {
 
 // -----
 
+//     CHECK:LetStmt {{.*}}
+//CHECK-NEXT:`-VariableDecl {{.*}} Name<dictionary> Type<Attr>
+//CHECK-NEXT:  `-CallExpr {{.*}} Type<Attr>
+//CHECK-NEXT:    `-DeclRefExpr {{.*}} Type<Constraint>
+//CHECK-NEXT:      `-UserConstraintDecl {{.*}} Name<__builtin_createDictionaryAttr_constraint> ResultType<Attr>
+//CHECK-NEXT:        `Results`
+//CHECK-NEXT:          `-VariableDecl {{.*}} Name<> Type<Attr>
+//CHECK-NEXT:            `Constraints`
+//CHECK-NEXT:              `-AttrConstraintDecl {{.*}}
+//CHECK-NEXT:ReturnStmt {{.*}}
+
+Constraint getEmptyDict() -> Attr {
+  let dictionary = {};
+  return dictionary;
+}
+
+// -----
+
+//     CHECK:LetStmt {{.*}}
+//CHECK-NEXT:`-VariableDecl {{.*}} Name<dictionary> Type<Attr>
+//CHECK-NEXT:  `-CallExpr {{.*}} Type<Attr>
+//CHECK-NEXT:    `-DeclRefExpr {{.*}} Type<Constraint>
+//CHECK-NEXT:      `-UserConstraintDecl {{.*}} Name<__builtin_addEntryToDictionaryAttr_constraint> ResultType<Attr>
+//     CHECK:    `Arguments`
+//CHECK-NEXT:      |-CallExpr {{.*}} Type<Attr>
+//CHECK-NEXT:      | `-DeclRefExpr {{.*}} Type<Constraint>
+//CHECK-NEXT:      |   `-UserConstraintDecl {{.*}} Name<__builtin_createDictionaryAttr_constraint> ResultType<Attr>
+//CHECK-NEXT:      |     `Results`
+//CHECK-NEXT:      |       `-VariableDecl {{.*}} Name<> Type<Attr>
+//CHECK-NEXT:      |         `Constraints`
+//CHECK-NEXT:      |           `-AttrConstraintDecl {{.*}}
+//CHECK-NEXT:      |-DeclRefExpr {{.*}} Type<Attr>
+//CHECK-NEXT:      | `-VariableDecl {{.*}} Name<dict0> Type<Attr>
+//CHECK-NEXT:      |   `-AttributeExpr {{.*}} Value<""test"">
+//CHECK-NEXT:      `-AttributeExpr {{.*}} Value<""String"">
+//CHECK-NEXT:ReturnStmt {{.*}}
+
+Constraint getPopulatedDict() -> Attr {
+  let dictionary = { test = "String" };
+  return dictionary;
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // CallExpr
 //===----------------------------------------------------------------------===//