[mlir][llvm] Add expect intrinsics.

[gysit]

Summary

The revision adds the LLVM expect and expect.with.probability
intrinsics.

[gysit]

Import the probability of the llvm.expect_with_probablility intrinsic into an attribute.

[Dinistro]

LGTM!

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

 def LLVM_USubWithOverflowOp : LLVM_ArithWithOverflowOp<"usub.with.overflow">;
 def LLVM_SMulWithOverflowOp : LLVM_ArithWithOverflowOp<"smul.with.overflow">;
 def LLVM_UMulWithOverflowOp : LLVM_ArithWithOverflowOp<"umul.with.overflow">;
 
 def LLVM_AssumeOp
   : LLVM_ZeroResultIntrOp<"assume", []>, Arguments<(ins I1:$cond)>;
 
 //
+// Expect intrinsics.
+//
+
+def LLVM_ExpectOp
+  : LLVM_OneResultIntrOp<"expect", [], [0],
+                         [Pure, SameOperandsAndResultType]> {
+  let arguments = (ins AnySignlessInteger:$val,
+                       AnySignlessInteger:$expected);
+  let assemblyFormat = "$val `,` $expected attr-dict `:` type($val)";
+}
+
+def LLVM_ExpectWithProbabilityOp
+  : LLVM_OneResultIntrOp<"expect.with.probability", [], [0],
+                         [Pure, AllTypesMatch<["val", "expected", "res"]>]> {
+  let arguments = (ins AnySignlessInteger:$val,
+                       AnySignlessInteger:$expected,
+                       F64Attr:$prob);
+  string llvmBuilder = [{
+    createIntrinsicCall(
+      builder, llvm::Intrinsic::expect_with_probability,
+      {$val, $expected, llvm::ConstantFP::get(builder.getDoubleTy(), $prob)},
+      {$_resultType});
+  }];
+  string mlirBuilder = [{
+    $res = $_builder.create<LLVM::ExpectWithProbabilityOp>(
+      $_location, $val, $expected, $_float_attr($prob));
+  }];
+  let assemblyFormat = "$val `,` $expected `,` $prob attr-dict `:` type($val)";
+}
+
+//
 // Coroutine intrinsics.
 //
 
 def LLVM_CoroIdOp : LLVM_IntrOp<"coro.id", [], [], [], 1> {
   let arguments = (ins I32:$align,
                        LLVM_i8Ptr:$promise,
                        LLVM_i8Ptr:$coroaddr,
                        LLVM_i8Ptr:$fnaddrs);

 
 def LLVM_Trap : LLVM_ZeroResultIntrOp<"trap">;
 
 def LLVM_DebugTrap : LLVM_ZeroResultIntrOp<"debugtrap">;
 
 def LLVM_UBSanTrap : LLVM_ZeroResultIntrOp<"ubsantrap"> {
   let arguments = (ins I8Attr:$failureKind);
   string llvmBuilder = [{
-    builder.CreateIntrinsic(
-      llvm::Intrinsic::ubsantrap, {}, {builder.getInt8($failureKind)});
+    createIntrinsicCall(
+      builder, llvm::Intrinsic::ubsantrap, {builder.getInt8($failureKind)});
   }];
   string mlirBuilder = [{
     $_op =
       $_builder.create<LLVM::UBSanTrap>($_location, $_int_attr($failureKind));
   }];
 }
 
 /// Create a call to vscale intrinsic.

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

class LLVM_OpBase<Dialect dialect, string mnemonic, list<Trait> traits = []> :
   // following $-variables exist:
   //   - $name - substituted by the remapped `inst` operand value at the index
   //             of the MLIR operation argument with the given name, or if the
   //             name matches the result name, by a reference to store the
   //             result of the newly created MLIR operation to;
   //   - $_op - substituted by a reference to store the newly created MLIR
   //            operation (only for MLIR operations that return no result);
   //   - $_int_attr - substituted by a call to an integer attribute matcher;
+  //   - $_float_attr - substituted by a call to a float attribute matcher;
   //   - $_var_attr - substituted by a call to a variable attribute matcher;
   //   - $_resultType - substituted with the MLIR result type;
   //   - $_location - substituted with the MLIR location;
   //   - $_builder - substituted with the MLIR builder;
   //   - $_qualCppClassName - substitiuted with the MLIR operation class name.
   // Always either store a reference to the result of the newly created
   // operation, or to the operation itself if it does not return a result.
   // Additionally, `$$` can be used to produce the dollar character.

mlir/include/mlir/Target/LLVMIR/ModuleImport.h

public:
 
   /// Converts a range of LLVM values to a range of MLIR values using the
   /// `convertValue` method, or returns failure if the conversion fails.
   FailureOr<SmallVector<Value>> convertValues(ArrayRef<llvm::Value *> values);
 
   /// Converts `value` to an integer attribute. Asserts if the matching fails.
   IntegerAttr matchIntegerAttr(llvm::Value *value);
 
+  /// Converts `value` to a float attribute. Asserts if the matching fails.
+  FloatAttr matchFloatAttr(llvm::Value *value);
+
   /// Converts `value` to a local variable attribute. Asserts if the matching
   /// fails.
   DILocalVariableAttr matchLocalVariableAttr(llvm::Value *value);
 
   /// Converts `value` to an array of symbol references pointing to alias scope
   /// operations, or returns failure if the conversion fails.
   FailureOr<SmallVector<SymbolRefAttr>>
   matchAliasScopeAttrs(llvm::Value *value);

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp

 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InlineAsm.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/MatrixBuilder.h"
 #include "llvm/IR/Operator.h"
 
 using namespace mlir;
 using namespace mlir::LLVM;
+using mlir::LLVM::detail::createIntrinsicCall;
 using mlir::LLVM::detail::getLLVMConstant;
 
 #include "mlir/Dialect/LLVMIR/LLVMConversionEnumsToLLVM.inc"
 
 static llvm::FastMathFlags getFastmathFlags(FastmathFlagsInterface &op) {
   using llvmFMF = llvm::FastMathFlags;
   using FuncT = void (llvmFMF::*)(bool);
   const std::pair<FastmathFlags, FuncT> handlers[] = {

mlir/lib/Target/LLVMIR/ModuleImport.cpp

ModuleImport::convertValues(ArrayRef<llvm::Value *> values) {
   return remapped;
 }
 
 IntegerAttr ModuleImport::matchIntegerAttr(llvm::Value *value) {
   IntegerAttr integerAttr;
   FailureOr<Value> converted = convertValue(value);
   bool success = succeeded(converted) &&
                  matchPattern(*converted, m_Constant(&integerAttr));
-  assert(success && "expected a constant value");
+  assert(success && "expected a constant integer value");
   (void)success;
   return integerAttr;
 }
 
+FloatAttr ModuleImport::matchFloatAttr(llvm::Value *value) {
+  FloatAttr floatAttr;
+  FailureOr<Value> converted = convertValue(value);
+  bool success =
+      succeeded(converted) && matchPattern(*converted, m_Constant(&floatAttr));
+  assert(success && "expected a constant float value");
+  (void)success;
+  return floatAttr;
+}
+
 DILocalVariableAttr ModuleImport::matchLocalVariableAttr(llvm::Value *value) {
   auto *nodeAsVal = cast<llvm::MetadataAsValue>(value);
   auto *node = cast<llvm::DILocalVariable>(nodeAsVal->getMetadata());
   return debugImporter->translate(node);
 }
 
 FailureOr<SmallVector<SymbolRefAttr>>
 ModuleImport::matchAliasScopeAttrs(llvm::Value *value) {

mlir/test/Target/LLVMIR/Import/intrinsic.ll

 ; CHECK-LABEL: @assume
 ; CHECK-SAME:  %[[TRUE:[a-zA-Z0-9]+]]
 define void @assume(i1 %true) {
   ; CHECK:  "llvm.intr.assume"(%[[TRUE]]) : (i1) -> ()
   call void @llvm.assume(i1 %true)
   ret void
 }
 
+; CHECK-LABEL: @expect
+; CHECK-SAME:  %[[VAL:[a-zA-Z0-9]+]]
+define void @expect(i32 %0) {
+  ; CHECK:  %[[EXP:.+]] = llvm.mlir.constant(42 : i32) : i32
+  ; CHECK:  llvm.intr.expect %[[VAL]], %[[EXP]] : i32
+  %2 = call i32 @llvm.expect.i32(i32 %0, i32 42)
+  ret void
+}
+
+; CHECK-LABEL: @expect_with_probability
+; CHECK-SAME:  %[[VAL:[a-zA-Z0-9]+]]
+define void @expect_with_probability(i16 %0) {
+  ; CHECK:  %[[EXP:.+]] = llvm.mlir.constant(42 : i16) : i16
+  ; CHECK:  llvm.intr.expect.with.probability %[[VAL]], %[[EXP]], 5.000000e-01 : i16
+  %2 = call i16 @llvm.expect.with.probability.i16(i16 %0, i16 42, double 0.5)
+  ret void
+}
+
 ; CHECK-LABEL:  llvm.func @coro_id
 define void @coro_id(i32 %0, ptr %1) {
   ; CHECK: llvm.intr.coro.id %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> !llvm.token
   %3 = call token @llvm.coro.id(i32 %0, ptr %1, ptr %1, ptr null)
   ret void
 }
 
 ; CHECK-LABEL:  llvm.func @coro_begin

 declare { i32, i1 } @llvm.ssub.with.overflow.i32(i32, i32)
 declare { <8 x i32>, <8 x i1> } @llvm.ssub.with.overflow.v8i32(<8 x i32>, <8 x i32>)
 declare { i32, i1 } @llvm.usub.with.overflow.i32(i32, i32)
 declare { <8 x i32>, <8 x i1> } @llvm.usub.with.overflow.v8i32(<8 x i32>, <8 x i32>)
 declare { i32, i1 } @llvm.smul.with.overflow.i32(i32, i32)
 declare { <8 x i32>, <8 x i1> } @llvm.smul.with.overflow.v8i32(<8 x i32>, <8 x i32>)
 declare { i32, i1 } @llvm.umul.with.overflow.i32(i32, i32)
 declare { <8 x i32>, <8 x i1> } @llvm.umul.with.overflow.v8i32(<8 x i32>, <8 x i32>)
+declare i32 @llvm.expect.i32(i32, i32)
+declare i16 @llvm.expect.with.probability.i16(i16, i16, double immarg)
 declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr)
 declare ptr @llvm.coro.begin(token, ptr writeonly)
 declare i64 @llvm.coro.size.i64()
 declare i32 @llvm.coro.size.i32()
 declare i64 @llvm.coro.align.i64()
 declare i32 @llvm.coro.align.i32()
 declare token @llvm.coro.save(ptr)
 declare i8 @llvm.coro.suspend(token, i1)

mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir

 llvm.func @umul_with_overflow_test(%arg0: i32, %arg1: i32, %arg2: vector<8xi32>, %arg3: vector<8xi32>) {
   // CHECK: call { i32, i1 } @llvm.umul.with.overflow.i32
   "llvm.intr.umul.with.overflow"(%arg0, %arg1) : (i32, i32) -> !llvm.struct<(i32, i1)>
   // CHECK: call { <8 x i32>, <8 x i1> } @llvm.umul.with.overflow.v8i32
   "llvm.intr.umul.with.overflow"(%arg2, %arg3) : (vector<8xi32>, vector<8xi32>) -> !llvm.struct<(vector<8xi32>, vector<8xi1>)>
   llvm.return
 }
 
+// CHECK-LABEL: @expect
+llvm.func @expect(%arg0: i32) {
+  %0 = llvm.mlir.constant(42 : i32) : i32
+  // CHECK: call i32 @llvm.expect.i32(i32 %{{.*}}, i32 42)
+  %1 = llvm.intr.expect %arg0, %0 : i32
+  llvm.return
+}
+
+// CHECK-LABEL: @expect_with_probability
+llvm.func @expect_with_probability(%arg0: i16) {
+  %0 = llvm.mlir.constant(42 : i16) : i16
+  // CHECK: call i16 @llvm.expect.with.probability.i16(i16 %{{.*}}, i16 42, double 5.000000e-01)
+  %1 = llvm.intr.expect.with.probability %arg0, %0, 5.000000e-01 : i16
+  llvm.return
+}
+
 // CHECK-LABEL: @coro_id
 llvm.func @coro_id(%arg0: i32, %arg1: !llvm.ptr<i8>) {
   // CHECK: call token @llvm.coro.id
   %null = llvm.mlir.null : !llvm.ptr<i8>
   llvm.intr.coro.id %arg0, %arg1, %arg1, %null : (i32, !llvm.ptr<i8>, !llvm.ptr<i8>, !llvm.ptr<i8>) -> !llvm.token
   llvm.return
 }
 

 // CHECK-DAG: declare { i32, i1 } @llvm.uadd.with.overflow.i32(i32, i32)
 // CHECK-DAG: declare { <8 x i32>, <8 x i1> } @llvm.uadd.with.overflow.v8i32(<8 x i32>, <8 x i32>)
 // CHECK-DAG: declare { i32, i1 } @llvm.ssub.with.overflow.i32(i32, i32)
 // CHECK-DAG: declare { <8 x i32>, <8 x i1> } @llvm.ssub.with.overflow.v8i32(<8 x i32>, <8 x i32>)
 // CHECK-DAG: declare { i32, i1 } @llvm.usub.with.overflow.i32(i32, i32)
 // CHECK-DAG: declare { <8 x i32>, <8 x i1> } @llvm.usub.with.overflow.v8i32(<8 x i32>, <8 x i32>)
 // CHECK-DAG: declare { i32, i1 } @llvm.umul.with.overflow.i32(i32, i32)
 // CHECK-DAG: declare { <8 x i32>, <8 x i1> } @llvm.umul.with.overflow.v8i32(<8 x i32>, <8 x i32>)
+// CHECK-DAG: declare i32 @llvm.expect.i32(i32, i32)
+// CHECK-DAG: declare i16 @llvm.expect.with.probability.i16(i16, i16, double immarg)
 // CHECK-DAG: declare token @llvm.coro.id(i32, ptr readnone, ptr nocapture readonly, ptr)
 // CHECK-DAG: declare ptr @llvm.coro.begin(token, ptr writeonly)
 // CHECK-DAG: declare i64 @llvm.coro.size.i64()
 // CHECK-DAG: declare i32 @llvm.coro.size.i32()
 // CHECK-DAG: declare token @llvm.coro.save(ptr)
 // CHECK-DAG: declare i8 @llvm.coro.suspend(token, i1)
 // CHECK-DAG: declare i1 @llvm.coro.end(ptr, i1)
 // CHECK-DAG: declare ptr @llvm.coro.free(token, ptr nocapture readonly)

mlir/tools/mlir-tblgen/LLVMIRConversionGen.cpp

while (StringLoc loc = findNextVariable(builderStrRef)) {
     } else if (isResultName(op, name)) {
       if (op.getNumResults() != 1)
         return emitError(record, "expected op to have one result");
       bs << "moduleImport.mapValue(inst)";
     } else if (name == "_op") {
       bs << "moduleImport.mapNoResultOp(inst)";
     } else if (name == "_int_attr") {
       bs << "moduleImport.matchIntegerAttr";
+    } else if (name == "_float_attr") {
+      bs << "moduleImport.matchFloatAttr";
     } else if (name == "_var_attr") {
       bs << "moduleImport.matchLocalVariableAttr";
     } else if (name == "_resultType") {
       bs << "moduleImport.convertType(inst->getType())";
     } else if (name == "_location") {
       bs << "moduleImport.translateLoc(inst->getDebugLoc())";
     } else if (name == "_builder") {
       bs << "odsBuilder";