[mlir] Add arith.addi_carry op

[kuhar]

Summary

The arith.addi_carry op implements integer addition with overflows. The carry is returned via the second result, as i1.

[antiagainst]

This patch actually contains both the new arith op definition and lowering to SPIR-V. What about splitting them into two patches for more targeted patches and a clean history?

[kuhar]

In D131893#3723852, @antiagainst wrote:

This patch actually contains both the new arith op definition and lowering to SPIR-V. What about splitting them into two patches for more targeted patches and a clean history?

Sure, it makes sense to split it.

[kuhar]

Removed conversion to spv.IAddCarry

[antiagainst]

Cool, just two small issues.

mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp

245

Use mlir::builder::getZeroAttr(carryTy) directly here?

mlir/test/Dialect/Arithmetic/canonicalize.mlir

549

Need to capture %arg0 behind a FileCheck symbol here.

[kuhar]

mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp

245

I'm not sure I can access the builder here, can I? This is a non-static method.

mlir/test/Dialect/Arithmetic/canonicalize.mlir

549

Other test cases in this file rely on the argument name being preserved. Do you think this is too fragile?

[antiagainst]

mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp

245

It's fine to construct a temporary mlir::Builder using the context with mlir::Builder(getContext()) and then call the getZeroAttr method. In MLIR attributes are uniqued and kept in the context and have a lifelime as a context. So it's fine. (mlir::OpBuilder is different---it can build operations and thus would need an insertion point and are stateful, etc.)

mlir/test/Dialect/Arithmetic/canonicalize.mlir

549

Normally we won't want to directly use the symbols from printer given they can be changed and thus break these tests. It's unlikely that we are gonna change how function arguments are printed though. So it's fine then, given other tests already use it.

[kuhar]

mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp

245

Ah, makes sense, thanks!

[bondhugula]

mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp

20

Not needed.

22

Not needed.

24

You don't need this include here.

228

No need of getSum(); just getType().

239–289

This is completely missing code comments. Please see another op's folding hook for reference.

268–269

Please add assert messages.

[kuhar]

Added comments, removed includes.

mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp

24

I removed all of these.
I'm used to a convention that roughly says that if you use something directly, it should be included in the same file, instead of relying on transitive includes that may change. Is there some relevant coding convention on includes documented somewhere?

239–289

I added some prose. Is this what you expected? Other fold functions in this file don't seem to explain much when it comes to the default case that calls the default folder (constFoldBinaryOp).

268–269

Actually I don't think these are needed/useful, as constFoldBinaryOp must already check the same things. I removed them.

[bondhugula]

mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp

24

Yes. https://llvm.org/docs/CodingStandards.html#minimal-list-of-includes

[bondhugula]

mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp

263

Typo: calculate

mlir/include/mlir/Dialect/Arithmetic/IR/ArithmeticOps.td

let description = [{
     // Tensor element-wise addition.
     %x = arith.addi %y, %z : tensor<4x?xi8>
     ```
   }];
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
+
+def Arith_AddICarryOp : Arith_Op<"addi_carry", [Commutative,
+    AllTypesMatch<["lhs", "rhs", "sum"]>]> {
+  let summary = "integer addition operation returning both the sum and carry";
+  let description = [{
+    The `addi_carry` operation takes two operands and returns two results: the
+    sum (same type as both operands), and the carry (boolean-like).
+
+    Example:
+
+    ```mlir
+    // Scalar addition.
+    %sum, %carry = arith.addi_carry %b, %c : i64, i1
+
+    // Vector element-wise addition.
+    %b:2 = arith.addi_carry %g, %h : vector<4xi32>, vector<4xi1>
+
+    // Tensor element-wise addition.
+    %c:2 = arith.addi_carry %y, %z : tensor<4x?xi8>, tensor<4x?xi1>
+    ```
+  }];
+
+  let arguments = (ins SignlessIntegerLike:$lhs, SignlessIntegerLike:$rhs);
+  let results = (outs SignlessIntegerLike:$sum, BoolLike:$carry);
+  let assemblyFormat = [{
+    $lhs `,` $rhs attr-dict `:` type($sum) `,` type($carry)
+  }];
+
+  let hasFolder = 1;
+
+  let extraClassDeclaration = [{
+    ::llvm::Optional<::llvm::SmallVector<int64_t, 4>> getShapeForUnroll();
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // SubIOp
 //===----------------------------------------------------------------------===//
 
 def Arith_SubIOp : Arith_IntBinaryOp<"subi"> {
   let summary = "integer subtraction operation";
   let hasFolder = 1;
   let hasCanonicalizer = 1;

mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp

 //===- ArithmeticOps.cpp - MLIR Arithmetic dialect ops implementation -----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
+#include <cassert>
 #include <utility>
 
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/CommonFolders.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
-#include "llvm/ADT/SmallString.h"
 
+#include "llvm/ADT/APInt.h"
 #include "llvm/ADT/APSInt.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/SmallVector.h"
 
 using namespace mlir;
 using namespace mlir::arith;
 
 //===----------------------------------------------------------------------===//
 // Pattern helpers
 //===----------------------------------------------------------------------===//
 

 
 void arith::AddIOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                                 MLIRContext *context) {
   patterns.add<AddIAddConstant, AddISubConstantRHS, AddISubConstantLHS>(
       context);
 }
 
 //===----------------------------------------------------------------------===//
+// AddICarryOp
+//===----------------------------------------------------------------------===//
+
+Optional<SmallVector<int64_t, 4>> arith::AddICarryOp::getShapeForUnroll() {
+  if (auto vt = getSum().getType().dyn_cast<VectorType>())
+    return llvm::to_vector<4>(vt.getShape());
+  return None;
+}
+
+// Returns the carry bit, assuming that `sum` is the result of addition of
+// `operand` and another number.
+static APInt calculateCarry(const APInt &sum, const APInt &operand) {
+  return sum.ult(operand) ? APInt::getAllOnes(1) : APInt::getZero(1);
+}
+
+LogicalResult arith::AddICarryOp::fold(ArrayRef<Attribute> operands,
+                                       SmallVectorImpl<OpFoldResult> &results) {
+  auto carryTy = getCarry().getType();
+  // addi_carry(x, 0) -> x, false
+  if (matchPattern(getRhs(), m_Zero())) {
+    auto carryZero = APInt::getZero(1);
+    Attribute falseValue;
+    if (operands[1].isa<IntegerAttr>())
+      falseValue = IntegerAttr::get(carryTy, carryZero);
+    else if (operands[1].isa<SplatElementsAttr>())
+      falseValue = SplatElementsAttr::get(carryTy, carryZero);
+    else
+      return failure();
+
+    results.push_back(getLhs());
+    results.push_back(falseValue);
+    return success();
+  }
+
+  if (Attribute sumAttr = constFoldBinaryOp<IntegerAttr>(
+          operands, [](APInt a, const APInt &b) { return std::move(a) + b; })) {
+    Attribute carryAttr;
+    if (auto lhs = operands[0].dyn_cast<IntegerAttr>()) {
+      auto sum = sumAttr.cast<IntegerAttr>();
+      carryAttr = IntegerAttr::get(
+          carryTy, calculateCarry(sum.getValue(), lhs.getValue()));
+    } else if (auto lhs = operands[0].dyn_cast<SplatElementsAttr>()) {
+      auto sum = sumAttr.cast<SplatElementsAttr>();
+      APInt carry = calculateCarry(sum.getSplatValue<APInt>(),
+                                   lhs.getSplatValue<APInt>());
+      carryAttr = SplatElementsAttr::get(carryTy, carry);
+    } else if (auto lhs = operands[0].dyn_cast<ElementsAttr>()) {
+      auto sum = sumAttr.cast<ElementsAttr>();
+      const int64_t numElems = sum.getNumElements();
+      assert(lhs.getNumElements() == numElems);
+      assert(numElems >= 0);
+      SmallVector<APInt> carryValues;
+      carryValues.reserve(static_cast<size_t>(numElems));
+
+      auto sumIt = sum.value_begin<APInt>();
+      auto lhsIt = lhs.value_begin<APInt>();
+      for (int64_t i = 0, e = numElems; i != e; ++i, ++sumIt, ++lhsIt)
+        carryValues.push_back(calculateCarry(*sumIt, *lhsIt));
+
+      carryAttr = DenseElementsAttr::get(carryTy, carryValues);
+    } else {
+      return failure();
+    }
+
+    results.push_back(sumAttr);
+    results.push_back(carryAttr);
+    return success();
+  }
+
+  return failure();
+}
+
+//===----------------------------------------------------------------------===//
 // SubIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::SubIOp::fold(ArrayRef<Attribute> operands) {
   // subi(x,x) -> 0
   if (getOperand(0) == getOperand(1))
     return Builder(getContext()).getZeroAttr(getType());
   // subi(x,0) -> x

mlir/test/Dialect/Arithmetic/canonicalize.mlir

 // CHECK-LABEL: @doubleAddSub2
 //  CHECK-NEXT:   return %arg0
 func.func @doubleAddSub2(%arg0: index, %arg1 : index) -> index {
   %sub = arith.subi %arg0, %arg1 : index
   %add = arith.addi %arg1, %sub : index
   return %add : index
 }
 
+// CHECK-LABEL: @addiCarryZeroRhs
+//  CHECK-NEXT:   %[[false:.+]] = arith.constant false
+//  CHECK-NEXT:   return %arg0, %[[false]]
+func.func @addiCarryZeroRhs(%arg0: i32) -> (i32, i1) {
+  %zero = arith.constant 0 : i32
+  %sum, %carry = arith.addi_carry %arg0, %zero: i32, i1
+  return %sum, %carry : i32, i1
+}
+
+// CHECK-LABEL: @addiCarryZeroRhsSplat
+//  CHECK-NEXT:   %[[false:.+]] = arith.constant dense<false> : vector<4xi1>
+//  CHECK-NEXT:   return %arg0, %[[false]]
+func.func @addiCarryZeroRhsSplat(%arg0: vector<4xi32>) -> (vector<4xi32>, vector<4xi1>) {
+  %zero = arith.constant dense<0> : vector<4xi32>
+  %sum, %carry = arith.addi_carry %arg0, %zero: vector<4xi32>, vector<4xi1>
+  return %sum, %carry : vector<4xi32>, vector<4xi1>
+}
+
+// CHECK-LABEL: @addiCarryZeroLhs
+//  CHECK-NEXT:   %[[false:.+]] = arith.constant false
+//  CHECK-NEXT:   return %arg0, %[[false]]
+func.func @addiCarryZeroLhs(%arg0: i32) -> (i32, i1) {
+  %zero = arith.constant 0 : i32
+  %sum, %carry = arith.addi_carry %zero, %arg0: i32, i1
+  return %sum, %carry : i32, i1
+}
+
+// CHECK-LABEL: @addiCarryConstants
+//  CHECK-DAG:    %[[false:.+]] = arith.constant false
+//  CHECK-DAG:    %[[c50:.+]] = arith.constant 50 : i32
+//  CHECK-NEXT:   return %[[c50]], %[[false]]
+func.func @addiCarryConstants() -> (i32, i1) {
+  %c13 = arith.constant 13 : i32
+  %c37 = arith.constant 37 : i32
+  %sum, %carry = arith.addi_carry %c13, %c37: i32, i1
+  return %sum, %carry : i32, i1
+}
+
+// CHECK-LABEL: @addiCarryConstantsOverflow1
+//  CHECK-DAG:    %[[true:.+]] = arith.constant true
+//  CHECK-DAG:    %[[c0:.+]] = arith.constant 0 : i32
+//  CHECK-NEXT:   return %[[c0]], %[[true]]
+func.func @addiCarryConstantsOverflow1() -> (i32, i1) {
+  %max = arith.constant 4294967295 : i32
+  %c1 = arith.constant 1 : i32
+  %sum, %carry = arith.addi_carry %max, %c1: i32, i1
+  return %sum, %carry : i32, i1
+}
+
+// CHECK-LABEL: @addiCarryConstantsOverflow2
+//  CHECK-DAG:    %[[true:.+]] = arith.constant true
+//  CHECK-DAG:    %[[c_2:.+]] = arith.constant -2 : i32
+// CHECK-NEXT:    return %[[c_2]], %[[true]]
+func.func @addiCarryConstantsOverflow2() -> (i32, i1) {
+  %max = arith.constant 4294967295 : i32
+  %sum, %carry = arith.addi_carry %max, %max: i32, i1
+  return %sum, %carry : i32, i1
+}
+
+// CHECK-LABEL: @addiCarryConstantsOverflowVector
+//  CHECK-DAG:    %[[sum:.+]] = arith.constant dense<[1, 6, 2, 14]> : vector<4xi32>
+//  CHECK-DAG:    %[[carry:.+]] = arith.constant dense<[false, false, true, false]> : vector<4xi1>
+// CHECK-NEXT:    return %[[sum]], %[[carry]]
+func.func @addiCarryConstantsOverflowVector() -> (vector<4xi32>, vector<4xi1>) {
+  %v1 = arith.constant dense<[1, 3, 3, 7]> : vector<4xi32>
+  %v2 = arith.constant dense<[0, 3, 4294967295, 7]> : vector<4xi32>
+  %sum, %carry = arith.addi_carry %v1, %v2 : vector<4xi32>, vector<4xi1>
+  return %sum, %carry : vector<4xi32>, vector<4xi1>
+}
+
+// CHECK-LABEL: @addiCarryConstantsSplatVector
+//   CHECK-DAG:   %[[sum:.+]] = arith.constant dense<3> : vector<4xi32>
+//   CHECK-DAG:   %[[carry:.+]] = arith.constant dense<false> : vector<4xi1>
+//  CHECK-NEXT:   return %[[sum]], %[[carry]]
+func.func @addiCarryConstantsSplatVector() -> (vector<4xi32>, vector<4xi1>) {
+  %v1 = arith.constant dense<1> : vector<4xi32>
+  %v2 = arith.constant dense<2> : vector<4xi32>
+  %sum, %carry = arith.addi_carry %v1, %v2 : vector<4xi32>, vector<4xi1>
+  return %sum, %carry : vector<4xi32>, vector<4xi1>
+}
+
 // CHECK-LABEL: @notCmpEQ
 //       CHECK:   %[[cres:.+]] = arith.cmpi ne, %arg0, %arg1 : i8
 //       CHECK:   return %[[cres]]
 func.func @notCmpEQ(%arg0: i8, %arg1: i8) -> i1 {
   %true = arith.constant true
   %cmp = arith.cmpi "eq", %arg0, %arg1 : i8
   %ncmp = arith.xori %cmp, %true : i1
   return %ncmp : i1

mlir/test/Dialect/Arithmetic/invalid.mlir

 func.func @func_with_ops(f32) {
 ^bb0(%a : f32):
   // expected-error@+1 {{'arith.addi' op operand #0 must be signless-integer-like}}
   %sf = arith.addi %a, %a : f32
 }
 
 // -----
 
+func.func @func_with_ops(%a: f32) {
+  // expected-error@+1 {{'arith.addi_carry' op operand #0 must be signless-integer-like}}
+  %r:2 = arith.addi_carry %a, %a : f32, i32
+  return
+}
+
+// -----
+
+func.func @func_with_ops(%a: i32) {
+  // expected-error@+1 {{'arith.addi_carry' op result #1 must be bool-like}}
+  %r:2 = arith.addi_carry %a, %a : i32, i32
+  return
+}
+
+// -----
+
+func.func @func_with_ops(%a: vector<8xi32>) {
+  // expected-error@+1 {{'arith.addi_carry' op if an operand is non-scalar, then all results must be non-scalar}}
+  %r:2 = arith.addi_carry %a, %a : vector<8xi32>, i1
+  return
+}
+
+// -----
+
+func.func @func_with_ops(%a: vector<8xi32>) {
+  // expected-error@+1 {{'arith.addi_carry' op all non-scalar operands/results must have the same shape and base type}}
+  %r:2 = arith.addi_carry %a, %a : vector<8xi32>, tensor<8xi1>
+  return
+}
+
+// -----
+
 func.func @func_with_ops(i32) {
 ^bb0(%a : i32):
   %sf = arith.addf %a, %a : i32  // expected-error {{'arith.addf' op operand #0 must be floating-point-like}}
 }
 
 // -----
 
 func.func @func_with_ops(i32) {

mlir/test/Dialect/Arithmetic/ops.mlir

 }
 
 // CHECK-LABEL: test_addi_scalable_vector
 func.func @test_addi_scalable_vector(%arg0 : vector<[8]xi64>, %arg1 : vector<[8]xi64>) -> vector<[8]xi64> {
   %0 = arith.addi %arg0, %arg1 : vector<[8]xi64>
   return %0 : vector<[8]xi64>
 }
 
+// CHECK-LABEL: test_addi_carry
+func.func @test_addi_carry(%arg0 : i64, %arg1 : i64) -> i64 {
+  %sum, %carry = arith.addi_carry %arg0, %arg1 : i64, i1
+  return %sum : i64
+}
+
+// CHECK-LABEL: test_addi_carry_tensor
+func.func @test_addi_carry_tensor(%arg0 : tensor<8x8xi64>, %arg1 : tensor<8x8xi64>) -> tensor<8x8xi64> {
+  %sum, %carry = arith.addi_carry %arg0, %arg1 : tensor<8x8xi64>, tensor<8x8xi1>
+  return %sum : tensor<8x8xi64>
+}
+
+// CHECK-LABEL: test_addi_carry_vector
+func.func @test_addi_carry_vector(%arg0 : vector<8xi64>, %arg1 : vector<8xi64>) -> vector<8xi64> {
+  %0:2 = arith.addi_carry %arg0, %arg1 : vector<8xi64>, vector<8xi1>
+  return %0#0 : vector<8xi64>
+}
+
+// CHECK-LABEL: test_addi_carry_scalable_vector
+func.func @test_addi_carry_scalable_vector(%arg0 : vector<[8]xi64>, %arg1 : vector<[8]xi64>) -> vector<[8]xi64> {
+  %0:2 = arith.addi_carry %arg0, %arg1 : vector<[8]xi64>, vector<[8]xi1>
+  return %0#0 : vector<[8]xi64>
+}
+
 // CHECK-LABEL: test_subi
 func.func @test_subi(%arg0 : i64, %arg1 : i64) -> i64 {
   %0 = arith.subi %arg0, %arg1 : i64
   return %0 : i64
 }
 
 // CHECK-LABEL: test_subi_tensor
 func.func @test_subi_tensor(%arg0 : tensor<8x8xi64>, %arg1 : tensor<8x8xi64>) -> tensor<8x8xi64> {