Fix linalg.dot over boolean tensors.

[jreiffers]

Summary

dot is currently miscompiled for booleans (uses add instead of or).

[jpienaar]

Could you add a test case here?

Also please update the description to make it specific to the error being addressed: XLA is a downstream project and the command line option is specific to there but not relevant here, if there were a GitHub issue you could reference it in body or reference this commit there post, but for commit title it should be standalone in this context.

[jreiffers]

Fix linalg.dot over boolean tensors.

For boolean tensors, this currently uses addi, which is inconsistent with numpy. Numpy uses saturating arithmetic.

>>> np.array([True]) + np.array([True])
array([ True])

[jreiffers]

In D129292#3635842, @jpienaar wrote:

<snip>

Done. Thanks.

[jreiffers]

Fix clang-format.

[bkramer]

Being compatible with numpy here seems fine to me.

[nicolasvasilache]

Thanks for fixing!

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

static MatchContractionResult isContractionInterfaceImpl(Operation *op) {
   if (linalgOp.getNumReductionLoops() == 0)
     return MatchContractionResult::NoReduction;
   if (llvm::any_of(mapRange,
                    [](AffineMap m) { return !m.isProjectedPermutation(); }))
     return MatchContractionResult::NotProjectedPermutations;
   // TODO: more fields than add/mul.
   if (!isAddMul<arith::AddFOp, arith::MulFOp>(linalgOp->getRegion(0).front()) &&
       !isAddMul<arith::AddIOp, arith::MulIOp>(linalgOp->getRegion(0).front()) &&
-      !isAddMul<complex::AddOp, complex::MulOp>(linalgOp->getRegion(0).front()))
+      !isAddMul<complex::AddOp, complex::MulOp>(
+          linalgOp->getRegion(0).front()) &&
+      !isAddMul<arith::OrIOp, arith::AndIOp>(linalgOp->getRegion(0).front()))
     return MatchContractionResult::NotAddMul;
   return MatchContractionResult::Success;
 }
 
 bool mlir::linalg::isaContractionOpInterface(LinalgOp linalgOp) {
   if (!linalgOp)
     return false;
   Operation *op = linalgOp.getOperation();

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

Value buildUnaryFn(UnaryFn unaryFn, Value arg) {
     llvm_unreachable("unsupported unary function");
   }
 
   // Build the binary functions defined by OpDSL.
   Value buildBinaryFn(BinaryFn binaryFn, Value arg0, Value arg1) {
     bool allComplex = isComplex(arg0) && isComplex(arg1);
     bool allFloatingPoint = isFloatingPoint(arg0) && isFloatingPoint(arg1);
     bool allInteger = isInteger(arg0) && isInteger(arg1);
+    bool allBool = allInteger && arg0.getType().getIntOrFloatBitWidth() == 1 &&
+                   arg1.getType().getIntOrFloatBitWidth() == 1;
     if (!allComplex && !allFloatingPoint && !allInteger)
       llvm_unreachable("unsupported non numeric type");
     OpBuilder builder = getBuilder();
     switch (binaryFn) {
     case BinaryFn::add:
       if (allComplex)
         return builder.create<complex::AddOp>(arg0.getLoc(), arg0, arg1);
       if (allFloatingPoint)
         return builder.create<arith::AddFOp>(arg0.getLoc(), arg0, arg1);
+      if (allBool)
+        return builder.create<arith::OrIOp>(arg0.getLoc(), arg0, arg1);
       return builder.create<arith::AddIOp>(arg0.getLoc(), arg0, arg1);
     case BinaryFn::sub:
       if (allComplex)
         return builder.create<complex::SubOp>(arg0.getLoc(), arg0, arg1);
       if (allFloatingPoint)
         return builder.create<arith::SubFOp>(arg0.getLoc(), arg0, arg1);
+      if (allBool)
+        llvm_unreachable("unsupported operation: sub with bools");
       return builder.create<arith::SubIOp>(arg0.getLoc(), arg0, arg1);
     case BinaryFn::mul:
       if (allComplex)
         return builder.create<complex::MulOp>(arg0.getLoc(), arg0, arg1);
       if (allFloatingPoint)
         return builder.create<arith::MulFOp>(arg0.getLoc(), arg0, arg1);
+      if (allBool)
+        return builder.create<arith::AndIOp>(arg0.getLoc(), arg0, arg1);
       return builder.create<arith::MulIOp>(arg0.getLoc(), arg0, arg1);
     case BinaryFn::max_signed:
       assert(!allComplex);
       if (allFloatingPoint)
         return builder.create<arith::MaxFOp>(arg0.getLoc(), arg0, arg1);
       return builder.create<arith::MaxSIOp>(arg0.getLoc(), arg0, arg1);
     case BinaryFn::min_signed:
       assert(!allComplex);

mlir/test/Dialect/Linalg/loops.mlir

 //   CHECKPARALLEL-DAG:   %[[a:.*]] = memref.load %[[A]][%{{.*}}] : memref<?xf32>
 //   CHECKPARALLEL-DAG:   %[[b:.*]] = memref.load %[[B]][%{{.*}}] : memref<?xf32>
 //   CHECKPARALLEL-DAG:   %[[inc:.*]] = arith.mulf %[[a]], %[[b]] : f32
 //   CHECKPARALLEL-DAG:   %[[c:.*]] = memref.load %[[C]][] : memref<f32>
 //   CHECKPARALLEL-DAG:   %[[res:.*]] = arith.addf %[[c]], %[[inc]] : f32
 //       CHECKPARALLEL:   store %[[res]], %[[C]][] : memref<f32>
 
 
+func.func @dot_int(%arg0: memref<?xi32>, %arg1: memref<?xi32>,
+                   %arg3: memref<i32>) {
+  // Verifies that we use the correct arith operations for integers.
+  linalg.dot ins(%arg0, %arg1 : memref<?xi32>, memref<?xi32>)
+             outs(%arg3 : memref<i32>)
+  return
+}
+// CHECK-LABEL: func @dot_int(
+//       CHECK:   %[[inc:.*]] = arith.muli {{.*}} : i32
+//  CHECK-NEXT:   %[[res:.*]] = arith.addi {{.*}}, %[[inc]] : i32
+//  CHECK-NEXT:   store %[[res]], {{.*}} : memref<i32>
+
+
+func.func @dot_bool(%arg0: memref<?xi1>, %arg1: memref<?xi1>,
+                    %arg3: memref<i1>) {
+  // Verifies that we use the correct (saturating) arith operations for booleans.
+  linalg.dot ins(%arg0, %arg1 : memref<?xi1>, memref<?xi1>)
+             outs(%arg3 : memref<i1>)
+  return
+}
+// CHECK-LABEL: func @dot_bool(
+//       CHECK:   %[[inc:.*]] = arith.andi {{.*}} : i1
+//  CHECK-NEXT:   %[[res:.*]] = arith.ori {{.*}}, %[[inc]] : i1
+//  CHECK-NEXT:   store %[[res]], {{.*}} : memref<i1>
+
+
 func.func @dot_view(%arg0: memref<?xf32, offset: ?, strides: [1]>, %arg1: memref<?xf32, offset: ?, strides: [1]>, %arg2: memref<f32>) {
   linalg.dot ins(%arg0, %arg1 : memref<?xf32, offset: ?, strides: [1]>,
                                 memref<?xf32, offset: ?, strides: [1]>)
             outs(%arg2:  memref<f32>)
   return
 }
 // CHECK-LABEL: func @dot_view(
 //       CHECK:   %{{.*}}: memref<?xf32, #[[$strided1D]]>, %{{.*}}: memref<?xf32, #[[$strided1D]]>, %{{.*}}: memref<f32>) {