[Mlir-commits] [mlir] 9ed9204 - [MLIR][Ploops] Add custom builders from ParallelOp and ReduceOp.

[Alexander Belyaev]

Author: Alexander Belyaev
Date: 2020-02-19T10:20:17+01:00
New Revision: 9ed920444f66a1796c952997d76977549384e201

URL: https://github.com/llvm/llvm-project/commit/9ed920444f66a1796c952997d76977549384e201
DIFF: https://github.com/llvm/llvm-project/commit/9ed920444f66a1796c952997d76977549384e201.diff

LOG: [MLIR][Ploops] Add custom builders from ParallelOp and ReduceOp.

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/LoopOps/LoopOps.td
    mlir/lib/Dialect/LoopOps/LoopOps.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/LoopOps/LoopOps.td b/mlir/include/mlir/Dialect/LoopOps/LoopOps.td
index 0d6ee5f3d830..67e9ed86d496 100644
--- a/mlir/include/mlir/Dialect/LoopOps/LoopOps.td
+++ b/mlir/include/mlir/Dialect/LoopOps/LoopOps.td
@@ -182,7 +182,10 @@ def ParallelOp : Loop_Op<"parallel",
   let builders = [
     OpBuilder<"Builder *builder, OperationState &result, "
               "ValueRange lowerBounds, ValueRange upperBounds, "
-              "ValueRange steps">
+              "ValueRange steps">,
+    OpBuilder<"Builder *builder, OperationState &result, "
+              "ValueRange lowerBounds, ValueRange upperBounds, "
+              "ValueRange steps, ArrayRef<Type> resultTypes">
   ];
 
   let extraClassDeclaration = [{
@@ -221,8 +224,8 @@ def ReduceOp : Loop_Op<"reduce", [HasParent<"ParallelOp">]> {
     Example:
 
     ```mlir
-      %zero = constant 0.0 : f32
-      loop.reduce(%zero) {
+      %operand = constant 1.0 : f32
+      loop.reduce(%operand) {
         ^bb0(%lhs : f32, %rhs: f32):
           %res = addf %lhs, %rhs : f32
           loop.reduce.return %res : f32
@@ -231,6 +234,12 @@ def ReduceOp : Loop_Op<"reduce", [HasParent<"ParallelOp">]> {
 
   }];
 
+  let skipDefaultBuilders = 1;
+  let builders = [
+    OpBuilder<"Builder *builder, OperationState &result, "
+              "Value operand">
+  ];
+
   let arguments = (ins AnyType:$operand);
   let regions = (region SizedRegion<1>:$reductionOperator);
 }

diff  --git a/mlir/lib/Dialect/LoopOps/LoopOps.cpp b/mlir/lib/Dialect/LoopOps/LoopOps.cpp
index 3550039b1423..3d2e89dc51c5 100644
--- a/mlir/lib/Dialect/LoopOps/LoopOps.cpp
+++ b/mlir/lib/Dialect/LoopOps/LoopOps.cpp
@@ -232,10 +232,17 @@ void ParallelOp::build(Builder *builder, OperationState &result, ValueRange lbs,
   result.addOperands(steps);
   Region *bodyRegion = result.addRegion();
   ParallelOp::ensureTerminator(*bodyRegion, *builder, result.location);
-  for (size_t i = 0; i < steps.size(); ++i)
+  for (size_t i = 0, e = steps.size(); i < e; ++i)
     bodyRegion->front().addArgument(builder->getIndexType());
 }
 
+void ParallelOp::build(Builder *builder, OperationState &result, ValueRange lbs,
+                       ValueRange ubs, ValueRange steps,
+                       ArrayRef<Type> resultTypes) {
+  result.addTypes(resultTypes);
+  build(builder, result, lbs, ubs, steps);
+}
+
 static LogicalResult verify(ParallelOp op) {
   // Check that there is at least one value in lowerBound, upperBound and step.
   // It is sufficient to test only step, because it is ensured already that the
@@ -354,6 +361,16 @@ ParallelOp mlir::loop::getParallelForInductionVarOwner(Value val) {
 // ReduceOp
 //===----------------------------------------------------------------------===//
 
+void ReduceOp::build(Builder *builder, OperationState &result, Value operand) {
+  auto type = operand.getType();
+  result.addOperands(operand);
+  Region *bodyRegion = result.addRegion();
+
+  Block *b = new Block();
+  b->addArguments(ArrayRef<Type>{type, type});
+  bodyRegion->getBlocks().insert(bodyRegion->end(), b);
+}
+
 static LogicalResult verify(ReduceOp op) {
   // The region of a ReduceOp has two arguments of the same type as its operand.
   auto type = op.operand().getType();