Added support for result groups

Author: Wheest

mlir/lib/AsmParser/Parser.cpp

@@ -1275,15 +1275,18 @@ OperationParser::parseSuccessors(SmallVectorImpl<Block *> &destinations) {
 /// Store the SSA names for the current operation as attrs for debug purposes.
 void OperationParser::storeSSANames(Operation *&op,
                                     ArrayRef<ResultRecord> resultIDs) {
-  if (op->getNumResults() > 1)
-    emitError("have not yet implemented support for multiple return values\n");
-
-  for (const ResultRecord &resIt : resultIDs) {
-    for (unsigned subRes : llvm::seq<unsigned>(0, std::get<1>(resIt))) {
-      op->setDiscardableAttr(
-          "mlir.ssaName",
-          StringAttr::get(getContext(), std::get<0>(resIt).drop_front(1)));
+
+  // Store the name(s) of the result(s) of this operation.
+  if (op->getNumResults() > 0) {
+    llvm::SmallVector<llvm::StringRef, 1> resultNames;
+    for (const ResultRecord &resIt : resultIDs) {
+      resultNames.push_back(std::get<0>(resIt).drop_front(1));
+      // Insert empty string for sub-results/result groups
+      for (unsigned int i = 1; i < std::get<1>(resIt); ++i)
+        resultNames.push_back(llvm::StringRef());
     }
+    op->setDiscardableAttr("mlir.resultNames",
+                           builder.getStrArrayAttr(resultNames));
   }
 
   // Store the name information of the arguments of this operation.

mlir/lib/IR/AsmPrinter.cpp

@@ -1300,7 +1300,8 @@ class SSANameState {
 
   /// Set the original identifier names if available. Used in debugging with
   /// `--retain-identifier-names`/`shouldRetainIdentifierNames` in ParserConfig
-  void setRetainedIdentifierNames(Operation &op);
+  void setRetainedIdentifierNames(Operation &op,
+                                  SmallVector<int, 2> &resultGroups);
 
   /// This is the value ID for each SSA value. If this returns NameSentinel,
   /// then the valueID has an entry in valueNames.
@@ -1573,7 +1574,7 @@ void SSANameState::numberValuesInOp(Operation &op) {
 
   // Set the original identifier names if available. Used in debugging with
   // `--retain-identifier-names`/`shouldRetainIdentifierNames` in ParserConfig
-  setRetainedIdentifierNames(op);
+  setRetainedIdentifierNames(op, resultGroups);
 
   unsigned numResults = op.getNumResults();
   if (numResults == 0) {
@@ -1597,25 +1598,33 @@ void SSANameState::numberValuesInOp(Operation &op) {
   }
 }
 
-void SSANameState::setRetainedIdentifierNames(Operation &op) {
-  // Get the original SSA for the result(s) if available
-  unsigned numResults = op.getNumResults();
-  if (numResults > 1)
-    llvm::outs()
-        << "have not yet implemented support for multiple return values\n";
-  else if (numResults == 1) {
-    Value resultBegin = op.getResult(0);
-    if (StringAttr ssaNameAttr = op.getAttrOfType<StringAttr>("mlir.ssaName")) {
-      setValueName(resultBegin, ssaNameAttr.strref());
-      op.removeDiscardableAttr("mlir.ssaName");
+void SSANameState::setRetainedIdentifierNames(
+    Operation &op, SmallVector<int, 2> &resultGroups) {
+  // Get the original names for the results if available
+  if (ArrayAttr resultNamesAttr =
+          op.getAttrOfType<ArrayAttr>("mlir.resultNames")) {
+    auto resultNames = resultNamesAttr.getValue();
+    auto results = op.getResults();
+    // Conservative in the case that the #results has changed
+    for (size_t i = 0; i < results.size() && i < resultNames.size(); ++i) {
+      auto resultName = resultNames[i].cast<StringAttr>().strref();
+      if (!resultName.empty()) {
+        if (!usedNames.count(resultName))
+          setValueName(results[i], resultName);
+        // If a result has a name, it is the start of a result group.
+        if (i > 0)
+          resultGroups.push_back(i);
+      }
     }
+    op.removeDiscardableAttr("mlir.resultNames");
   }
 
   // Get the original name for the op args if available
   if (ArrayAttr opArgNamesAttr =
           op.getAttrOfType<ArrayAttr>("mlir.opArgNames")) {
     auto opArgNames = opArgNamesAttr.getValue();
     auto opArgs = op.getOperands();
+    // Conservative in the case that the #operands has changed
     for (size_t i = 0; i < opArgs.size() && i < opArgNames.size(); ++i) {
       auto opArgName = opArgNames[i].cast<StringAttr>().strref();
       if (!usedNames.count(opArgName))
@@ -1636,6 +1645,7 @@ void SSANameState::setRetainedIdentifierNames(Operation &op) {
           op.getAttrOfType<ArrayAttr>("mlir.blockArgNames")) {
     auto blockArgNames = blockArgNamesAttr.getValue();
     auto blockArgs = op.getBlock()->getArguments();
+    // Conservative in the case that the #args has changed
     for (size_t i = 0; i < blockArgs.size() && i < blockArgNames.size(); ++i) {
       auto blockArgName = blockArgNames[i].cast<StringAttr>().strref();
       if (!usedNames.count(blockArgName))

mlir/test/IR/print-retain-identifiers.mlir

@@ -52,3 +52,41 @@ func.func @simple(i64, i1) -> i64 {
 }
 
 // -----
+
+//===----------------------------------------------------------------------===//
+// Test multiple return values
+//===----------------------------------------------------------------------===//
+
+func.func @select_min_max(%a: f64, %b: f64) -> (f64, f64) {
+  %gt = arith.cmpf "ogt", %a, %b : f64
+  // CHECK: %min, %max = scf.if %gt -> (f64, f64) {
+  %min, %max = scf.if %gt -> (f64, f64) {
+    scf.yield %b, %a : f64, f64
+  } else {
+    scf.yield %a, %b : f64, f64
+  }
+  // CHECK: return %min, %max : f64, f64
+  return %min, %max : f64, f64
+}
+
+// -----
+
+////===----------------------------------------------------------------------===//
+// Test multiple return values, with a grouped value tuple
+//===----------------------------------------------------------------------===//
+
+func.func @select_max(%a: f64, %b: f64, %c: f64, %d: f64) -> (f64, f64, f64, f64) {
+  // Find the max between %a and %b,
+  // with %c and %d being other values that are returned.
+  %gt = arith.cmpf "ogt", %a, %b : f64
+  // CHECK: %max, %others:2, %alt = scf.if %gt -> (f64, f64, f64, f64) {
+  %max, %others:2, %alt  = scf.if %gt -> (f64, f64, f64, f64) {
+    scf.yield %b, %a, %c, %d : f64, f64, f64, f64
+  } else {
+    scf.yield %a, %b, %d, %c : f64, f64, f64, f64
+  }
+  // CHECK: return %max, %others#0, %others#1, %alt : f64, f64, f64, f64
+  return %max, %others#0, %others#1, %alt : f64, f64, f64, f64
+}
+
+// -----