[mlir][LLVMIR] Ask ICmpOp to return vector<Nxi1> when needed

[myhsu]

Summary

If any of the operands for ICmpOp is a vector, returns a vector<Nxi1>, rather than an i1 type result.

[Mogball]

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

290

Can you define this builder in the source file?

[myhsu]

Put custom builder for ICmpOp into .cpp file.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

def LLVM_ICmpOp : LLVM_Op<"icmp", [NoSideEffect]> {
   let arguments = (ins ICmpPredicate:$predicate,
                    AnyTypeOf<[LLVM_ScalarOrVectorOf<AnyInteger>, LLVM_ScalarOrVectorOf<LLVM_AnyPointer>]>:$lhs,
                    AnyTypeOf<[LLVM_ScalarOrVectorOf<AnyInteger>, LLVM_ScalarOrVectorOf<LLVM_AnyPointer>]>:$rhs);
   let results = (outs LLVM_ScalarOrVectorOf<I1>:$res);
   let llvmBuilder = [{
     $res = builder.CreateICmp(getLLVMCmpPredicate($predicate), $lhs, $rhs);
   }];
   let builders = [
-    OpBuilder<(ins "ICmpPredicate":$predicate, "Value":$lhs, "Value":$rhs),
-    [{
-      build($_builder, $_state, IntegerType::get(lhs.getType().getContext(), 1),
-            predicate, lhs, rhs);
-    }]>];
+    OpBuilder<(ins "ICmpPredicate":$predicate, "Value":$lhs, "Value":$rhs)>
+  ];
   let hasCustomAssemblyFormat = 1;
 }
 
 // Predicate for float comparisons
 def FCmpPredicateFALSE  : I64EnumAttrCase<"_false", 0>;
 def FCmpPredicateOEQ    : I64EnumAttrCase<"oeq", 1>;
 def FCmpPredicateOGT    : I64EnumAttrCase<"ogt", 2>;
 def FCmpPredicateOGE    : I64EnumAttrCase<"oge", 3>;

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

if (!func)
     return op->emitOpError("'")
            << name << "' does not reference a valid LLVM function";
   if (func.isExternal())
     return op->emitOpError("'") << name << "' does not have a definition";
   return success();
 }
 
 //===----------------------------------------------------------------------===//
-// Printing/parsing for LLVM::CmpOp.
+// Printing, parsing and builder for LLVM::CmpOp.
 //===----------------------------------------------------------------------===//
 
+void ICmpOp::build(OpBuilder &builder, OperationState &result,
+                   ICmpPredicate predicate, Value lhs, Value rhs) {
+  auto boolType = IntegerType::get(lhs.getType().getContext(), 1);
+  if (LLVM::isCompatibleVectorType(lhs.getType()) ||
+      LLVM::isCompatibleVectorType(rhs.getType())) {
+    int64_t numLHSElements = 1, numRHSElements = 1;
+    if (LLVM::isCompatibleVectorType(lhs.getType()))
+      numLHSElements =
+          LLVM::getVectorNumElements(lhs.getType()).getFixedValue();
+    if (LLVM::isCompatibleVectorType(rhs.getType()))
+      numRHSElements =
+          LLVM::getVectorNumElements(rhs.getType()).getFixedValue();
+    build(builder, result,
+          VectorType::get({std::max(numLHSElements, numRHSElements)}, boolType),
+          predicate, lhs, rhs);
+  } else {
+    build(builder, result, boolType, predicate, lhs, rhs);
+  }
+}
+
 void ICmpOp::print(OpAsmPrinter &p) {
   p << " \"" << stringifyICmpPredicate(getPredicate()) << "\" " << getOperand(0)
     << ", " << getOperand(1);
   p.printOptionalAttrDict((*this)->getAttrs(), {"predicate"});
   p << " : " << getLhs().getType();
 }
 
 void FCmpOp::print(OpAsmPrinter &p) {

mlir/test/Dialect/LLVMIR/roundtrip.mlir

 // RUN: mlir-opt %s | mlir-opt | FileCheck %s
 
 // CHECK-LABEL: func @ops
 // CHECK-SAME: (%[[I32:.*]]: i32, %[[FLOAT:.*]]: f32, %[[I8PTR1:.*]]: !llvm.ptr<i8>, %[[I8PTR2:.*]]: !llvm.ptr<i8>, %[[BOOL:.*]]: i1, %[[VI8PTR1:.*]]: !llvm.vec<2 x ptr<i8>>)
 func.func @ops(%arg0: i32, %arg1: f32,
           %arg2: !llvm.ptr<i8>, %arg3: !llvm.ptr<i8>,
           %arg4: i1, %arg5 : !llvm.vec<2x!llvm.ptr<i8>>) {
 // Integer arithmetic binary operations.
 //
 // CHECK: {{.*}} = llvm.add %[[I32]], %[[I32]] : i32
 // CHECK: {{.*}} = llvm.sub %[[I32]], %[[I32]] : i32
 // CHECK: {{.*}} = llvm.mul %[[I32]], %[[I32]] : i32
 // CHECK: {{.*}} = llvm.udiv %[[I32]], %[[I32]] : i32
 // CHECK: {{.*}} = llvm.sdiv %[[I32]], %[[I32]] : i32
 // CHECK: {{.*}} = llvm.urem %[[I32]], %[[I32]] : i32
 // CHECK: {{.*}} = llvm.srem %[[I32]], %[[I32]] : i32
-// CHECK: {{.*}} = llvm.icmp "ne" %[[I32]], %[[I32]] : i32
-// CHECK: {{.*}} = llvm.icmp "ne" %[[I8PTR1]], %[[I8PTR1]] : !llvm.ptr<i8>
-// CHECK: {{.*}} = llvm.icmp "ne" %[[VI8PTR1]], %[[VI8PTR1]] : !llvm.vec<2 x ptr<i8>>
+// CHECK: %[[SCALAR_PRED0:.+]] = llvm.icmp "ne" %[[I32]], %[[I32]] : i32
+// CHECK: {{.*}} = llvm.add %[[SCALAR_PRED0]], %[[SCALAR_PRED0]] : i1
+// CHECK: %[[SCALAR_PRED1:.+]] = llvm.icmp "ne" %[[I8PTR1]], %[[I8PTR1]] : !llvm.ptr<i8>
+// CHECK: {{.*}} = llvm.add %[[SCALAR_PRED1]], %[[SCALAR_PRED1]] : i1
+// CHECK: %[[VEC_PRED:.+]] = llvm.icmp "ne" %[[VI8PTR1]], %[[VI8PTR1]] : !llvm.vec<2 x ptr<i8>>
+// CHECK: {{.*}} = llvm.add %[[VEC_PRED]], %[[VEC_PRED]] : vector<2xi1>
   %0 = llvm.add %arg0, %arg0 : i32
   %1 = llvm.sub %arg0, %arg0 : i32
   %2 = llvm.mul %arg0, %arg0 : i32
   %3 = llvm.udiv %arg0, %arg0 : i32
   %4 = llvm.sdiv %arg0, %arg0 : i32
   %5 = llvm.urem %arg0, %arg0 : i32
   %6 = llvm.srem %arg0, %arg0 : i32
   %7 = llvm.icmp "ne" %arg0, %arg0 : i32
+  %typecheck_7 = llvm.add %7, %7 : i1
   %ptrcmp = llvm.icmp "ne" %arg2, %arg2 : !llvm.ptr<i8>
+  %typecheck_ptrcmp = llvm.add %ptrcmp, %ptrcmp : i1
   %vptrcmp = llvm.icmp "ne" %arg5, %arg5 : !llvm.vec<2 x ptr<i8>>
+  %typecheck_vptrcmp = llvm.add %vptrcmp, %vptrcmp : vector<2 x i1>
 
 // Floating point binary operations.
 //
 // CHECK: {{.*}} = llvm.fadd %[[FLOAT]], %[[FLOAT]] : f32
 // CHECK: {{.*}} = llvm.fsub %[[FLOAT]], %[[FLOAT]] : f32
 // CHECK: {{.*}} = llvm.fmul %[[FLOAT]], %[[FLOAT]] : f32
 // CHECK: {{.*}} = llvm.fdiv %[[FLOAT]], %[[FLOAT]] : f32
 // CHECK: {{.*}} = llvm.frem %[[FLOAT]], %[[FLOAT]] : f32