[mlir] Fix support for lowering non-32-bit affine reductions.

[ftynse]

Summary

The existing implementation was always creating 32-bit constants for
floating-point and integer reductions regardless of the actual type, which
resulted in invalid IR being generated for any types other than f32 and i32
when lowering affine.parallel to SCF. Use the actual type instead.

[chelini]

Thanks, looks good to me.

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp

LogicalResult matchAndRewrite(AffineForOp op,
     rewriter.inlineRegionBefore(op.region(), scfForOp.region(),
                                 scfForOp.region().end());
     rewriter.replaceOp(op, scfForOp.results());
     return success();
   }
 };
 
 /// Returns the identity value associated with an AtomicRMWKind op.
-static Value getIdentityValue(AtomicRMWKind op, OpBuilder &builder,
-                              Location loc) {
+static Value getIdentityValue(AtomicRMWKind op, Type resultType,
+                              OpBuilder &builder, Location loc) {
   switch (op) {
   case AtomicRMWKind::addf:
-    return builder.create<ConstantOp>(loc, builder.getF32FloatAttr(0));
+    return builder.create<ConstantOp>(loc, builder.getFloatAttr(resultType, 0));
   case AtomicRMWKind::addi:
-    return builder.create<ConstantOp>(loc, builder.getI32IntegerAttr(0));
+    return builder.create<ConstantOp>(loc,
+                                      builder.getIntegerAttr(resultType, 0));
   case AtomicRMWKind::mulf:
-    return builder.create<ConstantOp>(loc, builder.getF32FloatAttr(1));
+    return builder.create<ConstantOp>(loc, builder.getFloatAttr(resultType, 1));
   case AtomicRMWKind::muli:
-    return builder.create<ConstantOp>(loc, builder.getI32IntegerAttr(1));
+    return builder.create<ConstantOp>(loc,
+                                      builder.getIntegerAttr(resultType, 1));
   // TODO: Add remaining reduction operations.
   default:
     (void)emitOptionalError(loc, "Reduction operation type not supported");
     break;
   }
   return nullptr;
 }
 
if (op.results().empty()) {
                                   parOp.region().end());
       rewriter.replaceOp(op, parOp.results());
       return success();
     }
     // Case with affine.parallel with reduction operations/return values.
     // scf.parallel handles the reduction operation differently unlike
     // affine.parallel.
     ArrayRef<Attribute> reductions = op.reductions().getValue();
-    for (Attribute reduction : reductions) {
+    for (auto pair : llvm::zip(reductions, op.getResultTypes())) {
       // For each of the reduction operations get the identity values for
       // initialization of the result values.
+      Attribute reduction = std::get<0>(pair);
+      Type resultType = std::get<1>(pair);
       Optional<AtomicRMWKind> reductionOp = symbolizeAtomicRMWKind(
           static_cast<uint64_t>(reduction.cast<IntegerAttr>().getInt()));
       assert(reductionOp.hasValue() &&
              "Reduction operation cannot be of None Type");
       AtomicRMWKind reductionOpValue = reductionOp.getValue();
-      identityVals.push_back(getIdentityValue(reductionOpValue, rewriter, loc));
+      identityVals.push_back(
+          getIdentityValue(reductionOpValue, resultType, rewriter, loc));
     }
     parOp = rewriter.create<scf::ParallelOp>(
         loc, lowerBoundTuple, upperBoundTuple, steps, identityVals,
         /*bodyBuilderFn=*/nullptr);
 
     //  Copy the body of the affine.parallel op.
     rewriter.eraseBlock(parOp.getBody());
     rewriter.inlineRegionBefore(op.region(), parOp.region(),

mlir/test/Conversion/AffineToStandard/lower-affine.mlir

 // CHECK-NEXT:      ^bb0(%[[LHS:.*]]: f32, %[[RHS:.*]]: f32):
 // CHECK-NEXT:        %[[RES:.*]] = mulf
 // CHECK-NEXT:        scf.reduce.return %[[RES]] : f32
 // CHECK-NEXT:      }
 // CHECK-NEXT:      scf.yield
 // CHECK-NEXT:    }
 // CHECK-NEXT:    return
 // CHECK-NEXT:  }
+
+/////////////////////////////////////////////////////////////////////
+
+func @affine_parallel_with_reductions_f64(%arg0: memref<3x3xf64>, %arg1: memref<3x3xf64>) -> (f64, f64) {
+  %0:2 = affine.parallel (%kx, %ky) = (0, 0) to (2, 2) reduce ("addf", "mulf") -> (f64, f64) {
+            %1 = affine.load %arg0[%kx, %ky] : memref<3x3xf64>
+            %2 = affine.load %arg1[%kx, %ky] : memref<3x3xf64>
+            %3 = mulf %1, %2 : f64
+            %4 = addf %1, %2 : f64
+            affine.yield %3, %4 : f64, f64
+          }
+  return %0#0, %0#1 : f64, f64
+}
+// CHECK-LABEL: @affine_parallel_with_reductions_f64
+// CHECK:  %[[LOWER_1:.*]] = constant 0 : index
+// CHECK:  %[[LOWER_2:.*]] = constant 0 : index
+// CHECK:  %[[UPPER_1:.*]] = constant 2 : index
+// CHECK:  %[[UPPER_2:.*]] = constant 2 : index
+// CHECK:  %[[STEP_1:.*]] = constant 1 : index
+// CHECK:  %[[STEP_2:.*]] = constant 1 : index
+// CHECK:  %[[INIT_1:.*]] = constant 0.000000e+00 : f64
+// CHECK:  %[[INIT_2:.*]] = constant 1.000000e+00 : f64
+// CHECK:  %[[RES:.*]] = scf.parallel (%[[I:.*]], %[[J:.*]]) = (%[[LOWER_1]], %[[LOWER_2]]) to (%[[UPPER_1]], %[[UPPER_2]]) step (%[[STEP_1]], %[[STEP_2]]) init (%[[INIT_1]], %[[INIT_2]]) -> (f64, f64) {
+// CHECK:    %[[VAL_1:.*]] = memref.load
+// CHECK:    %[[VAL_2:.*]] = memref.load
+// CHECK:    %[[PRODUCT:.*]] = mulf
+// CHECK:    %[[SUM:.*]] = addf
+// CHECK:    scf.reduce(%[[PRODUCT]]) : f64 {
+// CHECK:    ^bb0(%[[LHS:.*]]: f64, %[[RHS:.*]]: f64):
+// CHECK:      %[[RES:.*]] = addf
+// CHECK:      scf.reduce.return %[[RES]] : f64
+// CHECK:    }
+// CHECK:    scf.reduce(%[[SUM]]) : f64 {
+// CHECK:    ^bb0(%[[LHS:.*]]: f64, %[[RHS:.*]]: f64):
+// CHECK:      %[[RES:.*]] = mulf
+// CHECK:      scf.reduce.return %[[RES]] : f64
+// CHECK:    }
+// CHECK:    scf.yield
+// CHECK:  }
+
+/////////////////////////////////////////////////////////////////////
+
+func @affine_parallel_with_reductions_i64(%arg0: memref<3x3xi64>, %arg1: memref<3x3xi64>) -> (i64, i64) {
+  %0:2 = affine.parallel (%kx, %ky) = (0, 0) to (2, 2) reduce ("addi", "muli") -> (i64, i64) {
+            %1 = affine.load %arg0[%kx, %ky] : memref<3x3xi64>
+            %2 = affine.load %arg1[%kx, %ky] : memref<3x3xi64>
+            %3 = muli %1, %2 : i64
+            %4 = addi %1, %2 : i64
+            affine.yield %3, %4 : i64, i64
+          }
+  return %0#0, %0#1 : i64, i64
+}
+// CHECK-LABEL: @affine_parallel_with_reductions_i64
+// CHECK:  %[[LOWER_1:.*]] = constant 0 : index
+// CHECK:  %[[LOWER_2:.*]] = constant 0 : index
+// CHECK:  %[[UPPER_1:.*]] = constant 2 : index
+// CHECK:  %[[UPPER_2:.*]] = constant 2 : index
+// CHECK:  %[[STEP_1:.*]] = constant 1 : index
+// CHECK:  %[[STEP_2:.*]] = constant 1 : index
+// CHECK:  %[[INIT_1:.*]] = constant 0 : i64
+// CHECK:  %[[INIT_2:.*]] = constant 1 : i64
+// CHECK:  %[[RES:.*]] = scf.parallel (%[[I:.*]], %[[J:.*]]) = (%[[LOWER_1]], %[[LOWER_2]]) to (%[[UPPER_1]], %[[UPPER_2]]) step (%[[STEP_1]], %[[STEP_2]]) init (%[[INIT_1]], %[[INIT_2]]) -> (i64, i64) {
+// CHECK:    %[[VAL_1:.*]] = memref.load
+// CHECK:    %[[VAL_2:.*]] = memref.load
+// CHECK:    %[[PRODUCT:.*]] = muli
+// CHECK:    %[[SUM:.*]] = addi
+// CHECK:    scf.reduce(%[[PRODUCT]]) : i64 {
+// CHECK:    ^bb0(%[[LHS:.*]]: i64, %[[RHS:.*]]: i64):
+// CHECK:      %[[RES:.*]] = addi
+// CHECK:      scf.reduce.return %[[RES]] : i64
+// CHECK:    }
+// CHECK:    scf.reduce(%[[SUM]]) : i64 {
+// CHECK:    ^bb0(%[[LHS:.*]]: i64, %[[RHS:.*]]: i64):
+// CHECK:      %[[RES:.*]] = muli
+// CHECK:      scf.reduce.return %[[RES]] : i64
+// CHECK:    }
+// CHECK:    scf.yield
+// CHECK:  }