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

[herhut]

Thanks for adding the builders!

mlir/include/mlir/Dialect/LoopOps/LoopOps.td

240

Why is this called initValue? It is the value to be reduced.

mlir/lib/Dialect/LoopOps/LoopOps.cpp

239

Can you express the above builder by calling this builder? Or have the resultTypes have an empty default?

[rriddle]

mlir/lib/Dialect/LoopOps/LoopOps.cpp

248

Please cache the end value of this loop.

377

Nit: SmallVector -> ArrayRef?

[mehdi_amini]

mlir/include/mlir/Dialect/LoopOps/LoopOps.td

188

Isn't a loop always supposed to capture as operands the same types as the returned values? If so then why do we have a builder with resultTypes? I'd expect these to be inferred from the captured operands?

It is a problem with the current reduce, which does not model this, but that's a bug to fix right? And so adding a builder around the incorrect design does not seem right to me.

[pifon2a]

Addressed most comments.

[pifon2a]

mlir/include/mlir/Dialect/LoopOps/LoopOps.td

240

Then the documentation is quite confusing:

mlir
      %zero = constant 0.0 : f32
      loop.reduce(%zero) {
        ^bb0(%lhs : f32, %rhs: f32):
          %res = addf %lhs, %rhs : f32
          loop.reduce.return %res : f32
      } : f32

Looking at %zero I would assume that it is the initial value of the accumulator. I'll fix it.

[pifon2a]

Doc mod

[pifon2a]

mlir/include/mlir/Dialect/LoopOps/LoopOps.td

188

Ploops don't really capture any operands from which we can infer result types. The result types correspond to the ReduceOps in the body of the ploop.

[herhut]

mlir/include/mlir/Dialect/LoopOps/LoopOps.td

188

We have not added neutral values that can be used as initial values to parallel loops, yet, and I am not sure whether they should be mandatory. Having neutral values unlocks some additional reduction schemes but they are not always necessary. If we make them optional, we will need to provide return types to support that.
If we make them mandatory, we can adapt the builders and their use sites.

[mehdi_amini]

mlir/include/mlir/Dialect/LoopOps/LoopOps.td

188

I thought we discussed this and it would be added?
There is no guarantee that a loop would be executed once: as such it should take an input value for all returned to support this.

[nicolasvasilache]

mlir/lib/Dialect/LoopOps/LoopOps.cpp

239

I don't think this code path is tested?

370

I don't think this code path is tested?

[mehdi_amini]

Ping on ntv@ comments and my question?

[pifon2a]

@mehdi_amini @nicolasvasilache I am working on a PR that passes initial values to loop.parallel. After that PR this new builder won't be needed.

mlir/include/mlir/Dialect/LoopOps/LoopOps.td

let arguments = (ins Variadic<Index>:$lowerBound,
                        Variadic<Index>:$step);
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$region);
 
   let skipDefaultBuilders = 1;
   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 = [{
     Block *getBody() { return &region().front(); }
     iterator_range<Block::args_iterator> getInductionVars() {
       return {getBody()->args_begin(), getBody()->args_end()};
     }
   }];
```mlir
         ^bb0(%lhs : f32, %rhs: f32):
           %res = addf %lhs, %rhs : f32
           loop.reduce.return %res : f32
       } : f32
     ```
 
   }];
 
-  let arguments = (ins AnyType:$operand);
+  let skipDefaultBuilders = 1;
+  let builders = [
+    OpBuilder<"Builder *builder, OperationState &result, "
+              "Value initValue">
+  ];
+
+  let arguments = (ins AnyType:$initValue);
   let regions = (region SizedRegion<1>:$reductionOperator);
 }
 
 def ReduceReturnOp :
     Loop_Op<"reduce.return", [HasParent<"ReduceOp">, Terminator]> {
   let summary = "terminator for reduce operation";
   let description = [{
     "loop.reduce.return" is a special terminator operation for the block inside

mlir/lib/Dialect/LoopOps/LoopOps.cpp

void ParallelOp::build(Builder *builder, OperationState &result, ValueRange lbs,
   result.addOperands(ubs);
   result.addOperands(steps);
   Region *bodyRegion = result.addRegion();
   ParallelOp::ensureTerminator(*bodyRegion, *builder, result.location);
   for (size_t i = 0; i < steps.size(); ++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);
+  result.addOperands(lbs);
+  result.addOperands(ubs);
+  result.addOperands(steps);
+  Region *bodyRegion = result.addRegion();
+  ParallelOp::ensureTerminator(*bodyRegion, *builder, result.location);
+  for (size_t i = 0; i < steps.size(); ++i)
+    bodyRegion->front().addArgument(builder->getIndexType());
+}
+
 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
   // number of elements in lowerBound, upperBound and step are the same.
   Operation::operand_range stepValues = op.step();
   if (stepValues.empty())
     return op.emitOpError(
         "needs at least one tuple element for lowerBound, upperBound and step");
static LogicalResult verify(ParallelOp op) {
   if (op.results().size() != reductions.size())
     return op.emitOpError(
         "expects number of results to be the same as number of reductions");
 
   // Check that the types of the results and reductions are the same.
   for (auto resultAndReduce : llvm::zip(op.results(), reductions)) {
     auto resultType = std::get<0>(resultAndReduce).getType();
     auto reduceOp = std::get<1>(resultAndReduce);
-    auto reduceType = reduceOp.operand().getType();
+    auto reduceType = reduceOp.initValue().getType();
     if (resultType != reduceType)
       return reduceOp.emitOpError()
              << "expects type of reduce to be the same as result type: "
              << resultType;
   }
   return success();
 }
 
ParallelOp mlir::loop::getParallelForInductionVarOwner(Value val) {
   auto *containingInst = ivArg.getOwner()->getParentOp();
   return dyn_cast<ParallelOp>(containingInst);
 }
 
 //===----------------------------------------------------------------------===//
 // ReduceOp
 //===----------------------------------------------------------------------===//
 
+void ReduceOp::build(Builder *builder, OperationState &result,
+                     Value initValue) {
+  auto type = initValue.getType();
+  result.addOperands(initValue);
+  Region *bodyRegion = result.addRegion();
+
+  Block *b = new Block();
+  b->addArguments(SmallVector<Type, 2>({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();
+  auto type = op.initValue().getType();
   Block &block = op.reductionOperator().front();
   if (block.empty())
     return op.emitOpError("the block inside reduce should not be empty");
   if (block.getNumArguments() != 2 ||
       llvm::any_of(block.getArguments(), [&](const BlockArgument &arg) {
         return arg.getType() != type;
       }))
     return op.emitOpError()
static ParseResult parseReduceOp(OpAsmParser &parser, OperationState &result) {
   if (parser.parseColonType(resultType) ||
       parser.resolveOperand(operand, resultType, result.operands))
     return failure();
 
   return success();
 }
 
 static void print(OpAsmPrinter &p, ReduceOp op) {
-  p << op.getOperationName() << "(" << op.operand() << ") ";
+  p << op.getOperationName() << "(" << op.initValue() << ") ";
   p.printRegion(op.reductionOperator());
-  p << " : " << op.operand().getType();
+  p << " : " << op.initValue().getType();
 }
 
 //===----------------------------------------------------------------------===//
 // ReduceReturnOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(ReduceReturnOp op) {
   // The type of the return value should be the same type as the type of the
   // operand of the enclosing ReduceOp.
   auto reduceOp = cast<ReduceOp>(op.getParentOp());
-  Type reduceType = reduceOp.operand().getType();
+  Type reduceType = reduceOp.initValue().getType();
   if (reduceType != op.result().getType())
     return op.emitOpError() << "needs to have type " << reduceType
                             << " (the type of the enclosing ReduceOp)";
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/LoopOps/LoopOps.cpp.inc"