[Mlir-commits] [mlir] 08521ab - [mlir][EDSC] Allow conditionBuilder to capture the IfOp

[Nicolas Vasilache]

Author: Nicolas Vasilache
Date: 2020-07-17T11:16:26-04:00
New Revision: 08521abb3a7cafa50a0acdb16eca75162d96f514

URL: https://github.com/llvm/llvm-project/commit/08521abb3a7cafa50a0acdb16eca75162d96f514
DIFF: https://github.com/llvm/llvm-project/commit/08521abb3a7cafa50a0acdb16eca75162d96f514.diff

LOG: [mlir][EDSC] Allow conditionBuilder to capture the IfOp

When the IfOp returns values, it can easily be obtained from one of the Values.
However, when no values are returned, the information is lost.
This revision lets the caller specify a capture IfOp* to return the produced
IfOp.

Differential Revision: https://reviews.llvm.org/D84025

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/SCF/EDSC/Builders.h
    mlir/lib/Dialect/SCF/EDSC/Builders.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/SCF/EDSC/Builders.h b/mlir/include/mlir/Dialect/SCF/EDSC/Builders.h
index 1605f588bc00..50adec2f9b8b 100644
--- a/mlir/include/mlir/Dialect/SCF/EDSC/Builders.h
+++ b/mlir/include/mlir/Dialect/SCF/EDSC/Builders.h
@@ -36,12 +36,15 @@ scf::ValueVector loopNestBuilder(
 /// Adapters for building if conditions using the builder and the location
 /// stored in ScopedContext. 'thenBody' is mandatory, 'elseBody' can be omitted
 /// if the condition should not have an 'else' part.
-ValueRange
-conditionBuilder(TypeRange results, Value condition,
-                 function_ref<scf::ValueVector()> thenBody,
-                 function_ref<scf::ValueVector()> elseBody = nullptr);
+/// When `ifOp` is specified, the scf::IfOp is captured. This is particularly
+/// convenient for 0-result conditions.
+ValueRange conditionBuilder(TypeRange results, Value condition,
+                            function_ref<scf::ValueVector()> thenBody,
+                            function_ref<scf::ValueVector()> elseBody = nullptr,
+                            scf::IfOp *ifOp = nullptr);
 ValueRange conditionBuilder(Value condition, function_ref<void()> thenBody,
-                            function_ref<void()> elseBody = nullptr);
+                            function_ref<void()> elseBody = nullptr,
+                            scf::IfOp *ifOp = nullptr);
 
 } // namespace edsc
 } // namespace mlir

diff  --git a/mlir/lib/Dialect/SCF/EDSC/Builders.cpp b/mlir/lib/Dialect/SCF/EDSC/Builders.cpp
index 082c8c371ddc..2098ca1bf7d0 100644
--- a/mlir/lib/Dialect/SCF/EDSC/Builders.cpp
+++ b/mlir/lib/Dialect/SCF/EDSC/Builders.cpp
@@ -76,14 +76,17 @@ wrapIfBody(function_ref<scf::ValueVector()> body, TypeRange expectedTypes) {
 ValueRange
 mlir::edsc::conditionBuilder(TypeRange results, Value condition,
                              function_ref<scf::ValueVector()> thenBody,
-                             function_ref<scf::ValueVector()> elseBody) {
+                             function_ref<scf::ValueVector()> elseBody,
+                             scf::IfOp *ifOp) {
   assert(ScopedContext::getContext() && "EDSC ScopedContext not set up");
   assert(thenBody && "thenBody is mandatory");
 
-  auto ifOp = ScopedContext::getBuilderRef().create<scf::IfOp>(
+  auto newOp = ScopedContext::getBuilderRef().create<scf::IfOp>(
       ScopedContext::getLocation(), results, condition,
       wrapIfBody(thenBody, results), wrapIfBody(elseBody, results));
-  return ifOp.getResults();
+  if (ifOp)
+    *ifOp = newOp;
+  return newOp.getResults();
 }
 
 static std::function<void(OpBuilder &, Location)>
@@ -97,14 +100,17 @@ wrapZeroResultIfBody(function_ref<void()> body) {
 
 ValueRange mlir::edsc::conditionBuilder(Value condition,
                                         function_ref<void()> thenBody,
-                                        function_ref<void()> elseBody) {
+                                        function_ref<void()> elseBody,
+                                        scf::IfOp *ifOp) {
   assert(ScopedContext::getContext() && "EDSC ScopedContext not set up");
   assert(thenBody && "thenBody is mandatory");
 
-  ScopedContext::getBuilderRef().create<scf::IfOp>(
+  auto newOp = ScopedContext::getBuilderRef().create<scf::IfOp>(
       ScopedContext::getLocation(), condition, wrapZeroResultIfBody(thenBody),
       elseBody ? llvm::function_ref<void(OpBuilder &, Location)>(
                      wrapZeroResultIfBody(elseBody))
                : llvm::function_ref<void(OpBuilder &, Location)>(nullptr));
+  if (ifOp)
+    *ifOp = newOp;
   return {};
 }