[mlir][Vector] Support 0-D vectors in `CmpIOp`

[michalt]

Summary

Following the example of VectorOfAnyRankOf, I've done a few changes in the
.td files to help with adding the support for the 0-D case gradually.

[ftynse]

LGTM. Note that MLIR does not have any requirements for transition to be anyhow gradual. It does simplify the review, but that's all.

mlir/include/mlir/IR/OpBase.td

838

Nit: please consider adding a TODO: marker with "remove this when <X> is done" or a similar comment. Otherwise, this "temporary constraint" has big chances of becoming permanent :)

841

(I know this was not introduced in the present patch): one of the Of in VectorOfAnyRankOf sounds redundant.

[michalt]

Rebase and address comments

[michalt]

Thanks for the review! Could you land this? (I don't have commit access)

mlir/include/mlir/IR/OpBase.td

841

This sound ok to me, as in "vector of any rank, of integers". In any case, since this is temporary, I'd prefer to just leave it as is. CCing @nicolasvasilache in case he has an opinion :)

[nicolasvasilache]

mlir/include/mlir/IR/OpBase.td

841

The goal is to retire this and get back to simply VectorOf, the bigger the wart the simpler the cleanup?

[nicolasvasilache]

test is broken due to:

// Warning: these must be called in their textual order of definition in the
// file to not mess up FileCheck.

fixing and landing

[nicolasvasilache]

Fixing test.

[michalt]

In D115220#3187754, @nicolasvasilache wrote:

Fixing test.

Oops, sorry for that and thanks for fixing! :)

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

class Arith_CompareOp<string mnemonic, list<OpTrait> traits = []> :
     Arith_Op<mnemonic, traits # [SameTypeOperands, TypesMatchWith<
     "result type has i1 element type and same shape as operands",
     "lhs", "result", "::getI1SameShape($_self)">]> {
   let results = (outs BoolLike:$result);
 
   let assemblyFormat = "$predicate `,` $lhs `,` $rhs attr-dict `:` type($lhs)";
 }
 
+// Just like `Arith_CompareOp` but also admits 0-D vectors. Introduced
+// temporarily to allow gradual transition to 0-D vectors.
+class Arith_CompareOpOfAnyRank<string mnemonic, list<OpTrait> traits = []> :
+    Arith_CompareOp<mnemonic, traits> {
+  let results = (outs BoolLikeOfAnyRank:$result);
+}
+
 //===----------------------------------------------------------------------===//
 // ConstantOp
 //===----------------------------------------------------------------------===//
 
 def Arith_ConstantOp : Op<Arithmetic_Dialect, "constant",
     [ConstantLike, NoSideEffect,
      DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
      TypesMatchWith<

 
 //===----------------------------------------------------------------------===//
 // CeilDivUIOp
 //===----------------------------------------------------------------------===//
 
 def Arith_CeilDivUIOp : Arith_IntBinaryOp<"ceildivui"> {
   let summary = "unsigned ceil integer division operation";
   let description = [{
-    Unsigned integer division. Rounds towards positive infinity. Treats the 
+    Unsigned integer division. Rounds towards positive infinity. Treats the
     leading bit as the most significant, i.e. for `i16` given two's complement
     representation, `6 / -2 = 6 / (2^16 - 2) = 1`.
 
     Note: the semantics of division by zero is TBD; do NOT assume any specific
     behavior.
 
     Example:
 
def Arith_BitcastOp : Arith_CastOp<"bitcast", BitcastTypeConstraint,
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // CmpIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_CmpIOp : Arith_CompareOp<"cmpi"> {
+def Arith_CmpIOp : Arith_CompareOpOfAnyRank<"cmpi"> {
   let summary = "integer comparison operation";
   let description = [{
     The `cmpi` operation is a generic comparison for integer-like types. Its two
     arguments can be integers, vectors or tensors thereof as long as their types
     match. The operation produces an i1 for the former case, a vector or a
     tensor of i1 with the same shape as inputs in the other cases.
 
     Its first argument is an attribute that defines which type of comparison is
let description = [{
 
     // Generic form of the same operation.
     %x = "std.arith.cmpi"(%lhs, %rhs) {predicate = 0 : i64}
         : (vector<4xi64>, vector<4xi64>) -> vector<4xi1>
     ```
   }];
 
   let arguments = (ins Arith_CmpIPredicateAttr:$predicate,
-                       SignlessIntegerLike:$lhs,
-                       SignlessIntegerLike:$rhs);
+                       SignlessIntegerLikeOfAnyRank:$lhs,
+                       SignlessIntegerLikeOfAnyRank:$rhs);
 
   let builders = [
     OpBuilder<(ins "CmpIPredicate":$predicate, "Value":$lhs, "Value":$rhs), [{
       build($_builder, $_state, ::getI1SameShape(lhs.getType()),
             predicate, lhs, rhs);
     }]>
   ];
 

mlir/include/mlir/IR/OpBase.td

 //===----------------------------------------------------------------------===//
 
 // Type constraint for bool-like types: bools, vectors of bools, tensors of
 // bools.
 def BoolLike : TypeConstraint<Or<[I1.predicate, VectorOf<[I1]>.predicate,
                                   TensorOf<[I1]>.predicate]>,
     "bool-like">;
 
+// Temporary constraint to allow gradual transition to supporting 0-D vectors.
+def BoolLikeOfAnyRank : TypeConstraint<Or<[
+        I1.predicate,
+        VectorOfAnyRankOf<[I1]>.predicate,
+        TensorOf<[I1]>.predicate]>,
+    "bool-like">;
+
 // Type constraint for signless-integer-like types: signless integers, indices,
 // vectors of signless integers or indices, tensors of signless integers.
 def SignlessIntegerLike : TypeConstraint<Or<[
         AnySignlessIntegerOrIndex.predicate,
         VectorOf<[AnySignlessIntegerOrIndex]>.predicate,
         TensorOf<[AnySignlessIntegerOrIndex]>.predicate]>,
     "signless-integer-like">;
 
+// Temporary constraint to allow gradual transition to supporting 0-D vectors.
+def SignlessIntegerLikeOfAnyRank : TypeConstraint<Or<[
+        AnySignlessIntegerOrIndex.predicate,
+        VectorOfAnyRankOf<[AnySignlessIntegerOrIndex]>.predicate,
+        TensorOf<[AnySignlessIntegerOrIndex]>.predicate]>,
+    "signless-integer-like">;
+
 // Type constraint for float-like types: floats, vectors or tensors thereof.
 def FloatLike : TypeConstraint<Or<[AnyFloat.predicate,
         VectorOf<[AnyFloat]>.predicate, TensorOf<[AnyFloat]>.predicate]>,
     "floating-point-like">;
 
 // Type constraint for signless-integer-like or float-like types.
 def SignlessIntegerOrFloatLike : TypeConstraint<Or<[
     SignlessIntegerLike.predicate, FloatLike.predicate]>,

mlir/test/Conversion/ArithmeticToLLVM/arith-to-llvm.mlir

func @cmpf_2dvector(%arg0 : vector<4x3xf32>, %arg1 : vector<4x3xf32>) {
   // CHECK: %[[CMP:.*]] = llvm.fcmp "olt" %[[EXTRACT1]], %[[EXTRACT2]] : vector<3xf32>
   // CHECK: %[[INSERT:.*]] = llvm.insertvalue %[[CMP]], %2[0] : !llvm.array<4 x vector<3xi1>>
   %0 = arith.cmpf olt, %arg0, %arg1 : vector<4x3xf32>
   std.return
 }
 
 // -----
 
+// CHECK-LABEL: func @cmpi_0dvector(
+func @cmpi_0dvector(%arg0 : vector<i32>, %arg1 : vector<i32>) {
+  // CHECK: %[[ARG0:.*]] = builtin.unrealized_conversion_cast
+  // CHECK: %[[ARG1:.*]] = builtin.unrealized_conversion_cast
+  // CHECK: %[[CMP:.*]] = llvm.icmp "ult" %[[ARG0]], %[[ARG1]] : vector<1xi32>
+  %0 = arith.cmpi ult, %arg0, %arg1 : vector<i32>
+  std.return
+}
+
+// -----
+
 // CHECK-LABEL: func @cmpi_2dvector(
 func @cmpi_2dvector(%arg0 : vector<4x3xi32>, %arg1 : vector<4x3xi32>) {
   // CHECK: %[[ARG0:.*]] = builtin.unrealized_conversion_cast
   // CHECK: %[[ARG1:.*]] = builtin.unrealized_conversion_cast
   // CHECK: %[[EXTRACT1:.*]] = llvm.extractvalue %[[ARG0]][0] : !llvm.array<4 x vector<3xi32>>
   // CHECK: %[[EXTRACT2:.*]] = llvm.extractvalue %[[ARG1]][0] : !llvm.array<4 x vector<3xi32>>
   // CHECK: %[[CMP:.*]] = llvm.icmp "ult" %[[EXTRACT1]], %[[EXTRACT2]] : vector<3xi32>
   // CHECK: %[[INSERT:.*]] = llvm.insertvalue %[[CMP]], %2[0] : !llvm.array<4 x vector<3xi1>>
   %0 = arith.cmpi ult, %arg0, %arg1 : vector<4x3xi32>
   std.return
 }

mlir/test/Dialect/Arithmetic/ops.mlir

 }
 
 // CHECK-LABEL: test_cmpi_vector
 func @test_cmpi_vector(%arg0 : vector<8xi64>, %arg1 : vector<8xi64>) -> vector<8xi1> {
   %0 = arith.cmpi ult, %arg0, %arg1 : vector<8xi64>
   return %0 : vector<8xi1>
 }
 
+// CHECK-LABEL: test_cmpi_vector_0d
+func @test_cmpi_vector_0d(%arg0 : vector<i64>, %arg1 : vector<i64>) -> vector<i1> {
+  %0 = arith.cmpi ult, %arg0, %arg1 : vector<i64>
+  return %0 : vector<i1>
+}
+
 // CHECK-LABEL: test_cmpf
 func @test_cmpf(%arg0 : f64, %arg1 : f64) -> i1 {
   %0 = arith.cmpf oeq, %arg0, %arg1 : f64
   return %0 : i1
 }
 
 // CHECK-LABEL: test_cmpf_tensor
 func @test_cmpf_tensor(%arg0 : tensor<8x8xf64>, %arg1 : tensor<8x8xf64>) -> tensor<8x8xi1> {

mlir/test/Integration/Dialect/Vector/CPU/test-0-d-vectors.mlir

func @broadcast_0d(%a: f32) {
   vector.print %5: vector<2x1xf32>
 
   %6 = vector.broadcast %1 : vector<f32> to vector<2x3xf32>
   // CHECK: ( ( 42, 42, 42 ), ( 42, 42, 42 ) )
   vector.print %6: vector<2x3xf32>
   return
 }
 
+// CHECK_LABEL: arith_cmpi_0d
+func @arith_cmpi_0d(%smaller : vector<i32>, %bigger : vector<i32>) {
+  %0 = arith.cmpi ult, %smaller, %bigger : vector<i32>
+  // CHECK: ( 1 )
+  vector.print %0: vector<i1>
+
+  %1 = arith.cmpi ugt, %smaller, %bigger : vector<i32>
+  // CHECK: ( 0 )
+  vector.print %1: vector<i1>
+
+  %2 = arith.cmpi eq, %smaller, %bigger : vector<i32>
+  // CHECK: ( 0 )
+  vector.print %2: vector<i1>
+
+  return
+}
+
 func @entry() {
   %0 = arith.constant 42.0 : f32
   %1 = arith.constant dense<0.0> : vector<f32>
   %2 = call  @insert_element_0d(%0, %1) : (f32, vector<f32>) -> (vector<f32>)
   call  @extract_element_0d(%2) : (vector<f32>) -> ()
 
   %3 = arith.constant dense<42.0> : vector<f32>
   call  @print_vector_0d(%3) : (vector<f32>) -> ()
 
   %4 = arith.constant 42.0 : f32
   call  @splat_0d(%4) : (f32) -> ()
   call  @broadcast_0d(%4) : (f32) -> ()
 
+  %smaller = arith.constant dense<42> : vector<i32>
+  %bigger = arith.constant dense<4242> : vector<i32>
+  call  @arith_cmpi_0d(%smaller, %bigger) : (vector<i32>, vector<i32>) -> ()
+
   return
 }