Fix roundtrip of affine.if (#16)

Author: Frank Laub

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

@@ -322,8 +322,8 @@ def AffineIfOp : Affine_Op<"if",
     }
     ```
   }];
-  let results = (outs Variadic<AnyMemRef>:$results);
   let arguments = (ins Variadic<AnyType>);
+  let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$thenRegion, AnyRegion:$elseRegion);
 
   let skipDefaultBuilders = 1;
@@ -755,8 +755,7 @@ def AffineStoreOp : AffineStoreOpBase<"store", []> {
   let hasFolder = 1;
 }
 
-def AffineYieldOp : Affine_Op<"yield", [NoSideEffect, Terminator]>,
-    Arguments<(ins Variadic<AnyMemRef>:$values)> {
+def AffineYieldOp : Affine_Op<"yield", [NoSideEffect, Terminator]> {
   let summary = "Yield values to parent operation";
   let description = [{
     "affine.yield" yields zero or more SSA values from an affine op region and
@@ -770,7 +769,7 @@ def AffineYieldOp : Affine_Op<"yield", [NoSideEffect, Terminator]>,
     yielded.
     ```
   }];
-  let arguments = (ins Variadic<AnyMemRef>:$results);
+  let arguments = (ins Variadic<AnyType>:$results);
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result",
               [{ /* nothing to do */ }]>

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

@@ -1749,6 +1749,9 @@ static ParseResult parseAffineIfOp(OpAsmParser &parser,
         parser.getNameLoc(),
         "symbol operand count and integer set symbol count must match");
 
+  if (parser.parseOptionalArrowTypeList(result.types))
+    return failure();
+
   // Create the regions for 'then' and 'else'.  The latter must be created even
   // if it remains empty for the validity of the operation.
   result.regions.reserve(2);
@@ -1782,6 +1785,7 @@ static void print(OpAsmPrinter &p, AffineIfOp op) {
   p << "affine.if " << conditionAttr;
   printDimAndSymbolList(op.operand_begin(), op.operand_end(),
                         conditionAttr.getValue().getNumDims(), p);
+  p.printOptionalArrowTypeList(op.getResultTypes());
   p.printRegion(op.thenRegion(),
                 /*printEntryBlockArgs=*/false,
                 /*printBlockTerminators=*/op.getNumResults());

mlir/test/Dialect/Affine/ops.mlir

@@ -153,7 +153,7 @@ func @valid_symbol_affine_scope(%n : index, %A : memref<?xf32>) {
 
 // -----
 
-// CHECK-LABEL: @parallel
+// CHECK-LABEL: func @parallel
 // CHECK-SAME: (%[[N:.*]]: index)
 func @parallel(%N : index) {
   // CHECK: affine.parallel (%[[I0:.*]], %[[J0:.*]]) = (0, 0) to (symbol(%[[N]]), 100) step (10, 10)
@@ -166,3 +166,21 @@ func @parallel(%N : index) {
   }
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @affine_if
+func @affine_if() -> f32 {
+  // CHECK: %[[ZERO:.*]] = constant {{.*}} : f32
+  %zero = constant 0.0 : f32
+  // CHECK: %[[OUT:.*]] = affine.if {{.*}}() -> f32 {
+  %0 = affine.if affine_set<() : ()> () -> f32 {
+    // CHECK: affine.yield %[[ZERO]] : f32
+    affine.yield %zero : f32
+  } else {
+    // CHECK: affine.yield %[[ZERO]] : f32
+    affine.yield %zero : f32
+  }
+  // CHECK: return %[[OUT]] : f32
+  return %0 : f32
+}