[Mlir-commits] [mlir] [MLIR][NVVM] Add nvvm.fma Op (PR #184776)
Srinivasa Ravi
llvmlistbot at llvm.org
Thu Mar 5 04:32:55 PST 2026
https://github.com/Wolfram70 updated https://github.com/llvm/llvm-project/pull/184776
>From d6c9d9507564477baefd141a21e74d7855840e43 Mon Sep 17 00:00:00 2001
From: Srinivasa Ravi <srinivasar at nvidia.com>
Date: Thu, 5 Mar 2026 08:59:46 +0000
Subject: [PATCH 1/2] [MLIR][NVVM] Add nvvm.fma Op
Adds `nvvm.fma` Op to the NVVM dialect to perform fused multiply-add
operations.
PTX ISA Reference:
1. https://docs.nvidia.com/cuda/parallel-thread-execution/#floating-point-instructions-fma
2. https://docs.nvidia.com/cuda/parallel-thread-execution/#half-precision-floating-point-instructions-fma
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 31 ++
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 50 +++
.../Dialect/NVVM/NVVMToLLVMIRTranslation.cpp | 143 +++++++++
mlir/test/Target/LLVMIR/nvvm/fma/fma.mlir | 114 +++++++
.../Target/LLVMIR/nvvm/fma/fma_invalid.mlir | 89 ++++++
.../Target/LLVMIR/nvvm/fma/fma_vector.mlir | 294 ++++++++++++++++++
6 files changed, 721 insertions(+)
create mode 100644 mlir/test/Target/LLVMIR/nvvm/fma/fma.mlir
create mode 100644 mlir/test/Target/LLVMIR/nvvm/fma/fma_invalid.mlir
create mode 100644 mlir/test/Target/LLVMIR/nvvm/fma/fma_vector.mlir
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 43c7b3df73efe..7d98544c86bd4 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -6368,6 +6368,37 @@ def NVVM_SubFOp : NVVM_FloatBinaryOp<"subf"> {
let hasCanonicalizer = 1;
}
+def NVVM_FmaOp : NVVM_Op<"fma", [Pure, SameOperandsAndResultType]> {
+ let summary = [{
+ Performs floating point fused multiply-add operation with support for mixed
+ precision operands
+ }];
+ let description = [{
+ The `nvvm.fma` operation performs floating point fused multiply-add of three operands.
+ }];
+ let arguments = (ins
+ SIMTFloatType:$a,
+ SIMTFloatType:$b,
+ SIMTFloatType:$c,
+ FPArithRoundingMode:$rnd,
+ DefaultValuedAttr<SaturationModeSatOrNone, "SaturationMode::NONE">:$sat,
+ DefaultValuedAttr<BoolAttr, "false">:$ftz,
+ DefaultValuedAttr<BoolAttr, "false">:$relu,
+ DefaultValuedAttr<BoolAttr, "false">:$oob
+ );
+ let results = (outs SIMTFloatType:$res);
+ let assemblyFormat = "$a `,` $b `,` $c attr-dict `:` type($a)";
+ let hasVerifier = 1;
+
+ let extraClassDeclaration = [{
+ static void lowerFmaToLLVMIR(
+ Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder);
+ }];
+ let llvmBuilder = [{
+ NVVM::FmaOp::lowerFmaToLLVMIR(*op, moduleTranslation, builder);
+ }];
+}
+
//===----------------------------------------------------------------------===//
// NVVM tensormap.replace Op
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 354a02f4a2aae..2ab5f306bc24a 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3099,6 +3099,56 @@ LogicalResult NVVM::AddFOp::verify() {
return success();
}
+LogicalResult NVVM::FmaOp::verify() {
+ auto opType = getRes().getType();
+ mlir::NVVM::FPRoundingMode rndMode = getRnd();
+ mlir::NVVM::SaturationMode satMode = getSat();
+ bool isFTZ = getFtz();
+ bool isRelu = getRelu();
+ bool hasOOB = getOob();
+
+ auto getBaseFType = [](Type type) -> Type {
+ if (isa<VectorType>(type))
+ return cast<VectorType>(type).getElementType();
+ return type;
+ };
+
+ auto opBaseType = getBaseFType(opType);
+
+ if (rndMode == NVVM::FPRoundingMode::NONE)
+ return emitOpError("rounding mode must be specified");
+
+ if (isRelu && satMode == NVVM::SaturationMode::SAT)
+ return emitOpError("relu and saturation are not supported together");
+
+ if (hasOOB && (satMode == NVVM::SaturationMode::SAT || isFTZ))
+ return emitOpError("oob is not supported with saturation or FTZ");
+
+ if (!(opBaseType.isF16() || opBaseType.isBF16()) && (isRelu || hasOOB))
+ return emitOpError("relu and oob are only supported for f16 and bf16 fused "
+ "multiply-add operations");
+
+ if (opBaseType.isF64() && (satMode != NVVM::SaturationMode::NONE || isFTZ))
+ return emitOpError(
+ "FTZ and saturation are not supported for fused multiply-add "
+ "operations involving f64 type");
+
+ if (opBaseType.isF16() && rndMode != NVVM::FPRoundingMode::RN)
+ return emitOpError("only RN rounding mode is supported for f16 and "
+ "vector<2xf16> fused multiply-add operations");
+
+ if (opBaseType.isBF16()) {
+ if (rndMode != NVVM::FPRoundingMode::RN)
+ return emitOpError("only RN rounding mode is supported for bf16 and "
+ "vector<2xbf16> fused multiply-add operations");
+ if (satMode != NVVM::SaturationMode::NONE || isFTZ)
+ return emitOpError("FTZ and saturation are not supported for bf16 and "
+ "vector<2xbf16> fused multiply-add operations");
+ }
+
+ return success();
+}
+
/// Packs the given `field` into the `result`.
/// The `result` is 64-bits and each `field` can be 32-bits or narrower.
static llvm::Value *
diff --git a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
index 02067bb456b25..38111e7c0b707 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
@@ -563,6 +563,149 @@ void NVVM::AddFOp::lowerAddFToLLVMIR(Operation &op, LLVM::ModuleTranslation &mt,
}
}
+void NVVM::FmaOp::lowerFmaToLLVMIR(Operation &op, LLVM::ModuleTranslation &mt,
+ llvm::IRBuilderBase &builder) {
+ auto thisOp = cast<NVVM::FmaOp>(op);
+ llvm::SmallVector<llvm::Value *> args;
+ mlir::NVVM::FPRoundingMode rndMode = thisOp.getRnd();
+ unsigned rndIndex = static_cast<unsigned>(rndMode) - 1; // 1-4 mapped to 0-3
+ mlir::NVVM::SaturationMode satMode = thisOp.getSat();
+ bool isFTZ = thisOp.getFtz();
+ bool isRelu = thisOp.getRelu();
+ bool isSat = satMode == NVVM::SaturationMode::SAT;
+ bool isOOB = thisOp.getOob();
+
+ mlir::Type opType = thisOp.getRes().getType();
+ llvm::Type *opTypeLLVM = mt.convertType(opType);
+ bool isVectorAdd = opTypeLLVM->isVectorTy();
+
+ llvm::Value *argA = mt.lookupValue(thisOp.getA());
+ llvm::Value *argB = mt.lookupValue(thisOp.getB());
+ llvm::Value *argC = mt.lookupValue(thisOp.getC());
+
+ static constexpr llvm::Intrinsic::ID f16IDs[] = {
+ llvm::Intrinsic::nvvm_fma_rn_f16,
+ llvm::Intrinsic::nvvm_fma_rn_f16x2,
+ llvm::Intrinsic::nvvm_fma_rn_ftz_f16,
+ llvm::Intrinsic::nvvm_fma_rn_ftz_f16x2,
+ llvm::Intrinsic::nvvm_fma_rn_sat_f16,
+ llvm::Intrinsic::nvvm_fma_rn_sat_f16x2,
+ llvm::Intrinsic::nvvm_fma_rn_ftz_sat_f16,
+ llvm::Intrinsic::nvvm_fma_rn_ftz_sat_f16x2,
+ llvm::Intrinsic::nvvm_fma_rn_relu_f16,
+ llvm::Intrinsic::nvvm_fma_rn_relu_f16x2,
+ llvm::Intrinsic::nvvm_fma_rn_ftz_relu_f16,
+ llvm::Intrinsic::nvvm_fma_rn_ftz_relu_f16x2};
+
+ static constexpr llvm::Intrinsic::ID bf16IDs[] = {
+ llvm::Intrinsic::nvvm_fma_rn_bf16, llvm::Intrinsic::nvvm_fma_rn_bf16x2,
+ llvm::Intrinsic::nvvm_fma_rn_relu_bf16,
+ llvm::Intrinsic::nvvm_fma_rn_relu_bf16x2};
+
+ static constexpr llvm::Intrinsic::ID f32IDs[] = {
+ llvm::Intrinsic::nvvm_fma_rn_f,
+ llvm::Intrinsic::nvvm_fma_rm_f,
+ llvm::Intrinsic::nvvm_fma_rp_f,
+ llvm::Intrinsic::nvvm_fma_rz_f,
+ llvm::Intrinsic::nvvm_fma_rn_sat_f,
+ llvm::Intrinsic::nvvm_fma_rm_sat_f,
+ llvm::Intrinsic::nvvm_fma_rp_sat_f,
+ llvm::Intrinsic::nvvm_fma_rz_sat_f,
+ llvm::Intrinsic::nvvm_fma_rn_ftz_f,
+ llvm::Intrinsic::nvvm_fma_rm_ftz_f,
+ llvm::Intrinsic::nvvm_fma_rp_ftz_f,
+ llvm::Intrinsic::nvvm_fma_rz_ftz_f,
+ llvm::Intrinsic::nvvm_fma_rn_ftz_sat_f,
+ llvm::Intrinsic::nvvm_fma_rm_ftz_sat_f,
+ llvm::Intrinsic::nvvm_fma_rp_ftz_sat_f,
+ llvm::Intrinsic::nvvm_fma_rz_ftz_sat_f,
+ };
+
+ static constexpr llvm::Intrinsic::ID f64IDs[] = {
+ llvm::Intrinsic::nvvm_fma_rn_d, llvm::Intrinsic::nvvm_fma_rm_d,
+ llvm::Intrinsic::nvvm_fma_rp_d, llvm::Intrinsic::nvvm_fma_rz_d};
+
+ auto fmaIntrinsic = [&](llvm::Intrinsic::ID IID) -> llvm::Value * {
+ auto createFmaIntrinsicCall = [&](llvm::Intrinsic::ID IID, llvm::Value *a,
+ llvm::Value *b,
+ llvm::Value *c) -> llvm::CallInst * {
+ llvm::SmallVector<llvm::Value *, 3> callArgs;
+ callArgs.push_back(a);
+ callArgs.push_back(b);
+ callArgs.push_back(c);
+ return createIntrinsicCall(builder, IID, opTypeLLVM, callArgs);
+ };
+
+ if (isVectorAdd && (opTypeLLVM->getScalarType()->isFloatTy() ||
+ opTypeLLVM->getScalarType()->isDoubleTy())) {
+ llvm::Value *result = llvm::PoisonValue::get(
+ llvm::FixedVectorType::get(opTypeLLVM->getScalarType(), 2));
+ for (int64_t i = 0; i < 2; ++i) {
+ llvm::Value *argAElemi =
+ builder.CreateExtractElement(argA, builder.getInt32(i));
+ llvm::Value *argBElemi =
+ builder.CreateExtractElement(argB, builder.getInt32(i));
+ llvm::Value *argCElemi =
+ builder.CreateExtractElement(argC, builder.getInt32(i));
+ llvm::Value *sum =
+ createFmaIntrinsicCall(IID, argAElemi, argBElemi, argCElemi);
+ result = builder.CreateInsertElement(result, sum, builder.getInt32(i));
+ };
+ return result;
+ }
+
+ return createFmaIntrinsicCall(IID, argA, argB, argC);
+ }; // fmaIntrinsic end
+
+ // f16 + f16 -> f16 / vector<2xf16> + vector<2xf16> -> vector<2xf16>
+ // FIXME: Allow lowering to add.rn.ftz.f16x2 and add.rn.ftz.f16 here when the
+ // intrinsics are available.
+ if (opTypeLLVM->getScalarType()->isHalfTy()) {
+ llvm::Value *result;
+ if (isOOB) {
+ result = fmaIntrinsic(isRelu ? llvm::Intrinsic::nvvm_fma_rn_oob_relu
+ : llvm::Intrinsic::nvvm_fma_rn_oob);
+ } else {
+ unsigned index =
+ (isRelu << 3) | (isSat << 2) | (isFTZ << 1) |
+ isVectorAdd; // Op verifier ensures that this index is valid
+ result = fmaIntrinsic(f16IDs[index]);
+ }
+ mt.mapValue(thisOp.getRes(), result);
+ return;
+ }
+
+ // bf16 + bf16 -> bf16 / vector<2xbf16> + vector<2xbf16> -> vector<2xbf16>
+ if (opTypeLLVM->getScalarType()->isBFloatTy()) {
+ llvm::Value *result;
+ if (isOOB) {
+ result = fmaIntrinsic(isRelu ? llvm::Intrinsic::nvvm_fma_rn_oob_relu
+ : llvm::Intrinsic::nvvm_fma_rn_oob);
+ } else {
+ unsigned index =
+ (isRelu << 1) |
+ isVectorAdd; // Op verifier ensures that this index is valid
+ result = fmaIntrinsic(bf16IDs[index]);
+ }
+ mt.mapValue(thisOp.getRes(), result);
+ return;
+ }
+
+ // f64 + f64 -> f64 / vector<2xf64> + vector<2xf64> -> vector<2xf64>
+ if (opTypeLLVM->getScalarType()->isDoubleTy()) {
+ mt.mapValue(thisOp.getRes(), fmaIntrinsic(f64IDs[rndIndex]));
+ return;
+ }
+
+ // f32 + f32 -> f32 / vector<2xf32> + vector<2xf32> -> vector<2xf32>
+ const unsigned numRndModes = 4; // RN, RM, RP, RZ
+ if (opTypeLLVM->getScalarType()->isFloatTy()) {
+ unsigned index = ((isFTZ << 1) | isSat) * numRndModes + rndIndex;
+ mt.mapValue(thisOp.getRes(), fmaIntrinsic(f32IDs[index]));
+ return;
+ }
+}
+
namespace {
/// Implementation of the dialect interface that converts operations belonging
/// to the NVVM dialect to LLVM IR.
diff --git a/mlir/test/Target/LLVMIR/nvvm/fma/fma.mlir b/mlir/test/Target/LLVMIR/nvvm/fma/fma.mlir
new file mode 100644
index 0000000000000..236175daff21e
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/fma/fma.mlir
@@ -0,0 +1,114 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+llvm.func @fma_f16(%a: f16, %b: f16, %c: f16) -> f16 {
+ // CHECK-LABEL: define half @fma_f16(half %0, half %1, half %2) {
+ // CHECK-NEXT: %4 = call half @llvm.nvvm.fma.rn.f16(half %0, half %1, half %2)
+ // CHECK-NEXT: %5 = call half @llvm.nvvm.fma.rn.ftz.f16(half %0, half %1, half %4)
+ // CHECK-NEXT: %6 = call half @llvm.nvvm.fma.rn.sat.f16(half %0, half %1, half %5)
+ // CHECK-NEXT: %7 = call half @llvm.nvvm.fma.rn.ftz.sat.f16(half %0, half %1, half %6)
+ // CHECK-NEXT: %8 = call half @llvm.nvvm.fma.rn.relu.f16(half %0, half %1, half %7)
+ // CHECK-NEXT: %9 = call half @llvm.nvvm.fma.rn.ftz.relu.f16(half %0, half %1, half %8)
+ // CHECK-NEXT: %10 = call half @llvm.nvvm.fma.rn.oob.f16(half %0, half %1, half %9)
+ // CHECK-NEXT: %11 = call half @llvm.nvvm.fma.rn.oob.relu.f16(half %0, half %1, half %10)
+ // CHECK-NEXT: ret half %11
+ // CHECK-NEXT: }
+ %f0 = nvvm.fma %a, %b, %c {rnd = #nvvm.fp_rnd_mode<rn>} : f16
+ %f1 = nvvm.fma %a, %b, %f0 {rnd = #nvvm.fp_rnd_mode<rn>, ftz = true} : f16
+ %f2 = nvvm.fma %a, %b, %f1 {rnd = #nvvm.fp_rnd_mode<rn>, sat = #nvvm.sat_mode<sat>} : f16
+ %f3 = nvvm.fma %a, %b, %f2 {rnd = #nvvm.fp_rnd_mode<rn>, sat = #nvvm.sat_mode<sat>, ftz = true} : f16
+ %f4 = nvvm.fma %a, %b, %f3 {rnd = #nvvm.fp_rnd_mode<rn>, relu = true} : f16
+ %f5 = nvvm.fma %a, %b, %f4 {rnd = #nvvm.fp_rnd_mode<rn>, relu = true, ftz = true} : f16
+ %f6 = nvvm.fma %a, %b, %f5 {rnd = #nvvm.fp_rnd_mode<rn>, oob = true} : f16
+ %f7 = nvvm.fma %a, %b, %f6 {rnd = #nvvm.fp_rnd_mode<rn>, oob = true, relu = true} : f16
+ llvm.return %f7 : f16
+}
+
+llvm.func @fma_bf16(%a: bf16, %b: bf16, %c: bf16) -> bf16 {
+ // CHECK-LABEL: define bfloat @fma_bf16(bfloat %0, bfloat %1, bfloat %2) {
+ // CHECK-NEXT: %4 = call bfloat @llvm.nvvm.fma.rn.bf16(bfloat %0, bfloat %1, bfloat %2)
+ // CHECK-NEXT: %5 = call bfloat @llvm.nvvm.fma.rn.relu.bf16(bfloat %0, bfloat %1, bfloat %4)
+ // CHECK-NEXT: %6 = call bfloat @llvm.nvvm.fma.rn.oob.bf16(bfloat %0, bfloat %1, bfloat %5)
+ // CHECK-NEXT: %7 = call bfloat @llvm.nvvm.fma.rn.oob.relu.bf16(bfloat %0, bfloat %1, bfloat %6)
+ // CHECK-NEXT: ret bfloat %7
+ // CHECK-NEXT: }
+ %f0 = nvvm.fma %a, %b, %c {rnd = #nvvm.fp_rnd_mode<rn>} : bf16
+ %f1 = nvvm.fma %a, %b, %f0 {rnd = #nvvm.fp_rnd_mode<rn>, relu = true} : bf16
+ %f2 = nvvm.fma %a, %b, %f1 {rnd = #nvvm.fp_rnd_mode<rn>, oob = true} : bf16
+ %f3 = nvvm.fma %a, %b, %f2 {rnd = #nvvm.fp_rnd_mode<rn>, oob = true, relu = true} : bf16
+ llvm.return %f3 : bf16
+}
+
+llvm.func @fma_f32_rn(%a: f32, %b: f32, %c: f32) -> f32 {
+ // CHECK-LABEL: define float @fma_f32_rn(float %0, float %1, float %2) {
+ // CHECK-NEXT: %4 = call float @llvm.nvvm.fma.rn.f(float %0, float %1, float %2)
+ // CHECK-NEXT: %5 = call float @llvm.nvvm.fma.rn.ftz.f(float %0, float %1, float %4)
+ // CHECK-NEXT: %6 = call float @llvm.nvvm.fma.rn.sat.f(float %0, float %1, float %5)
+ // CHECK-NEXT: %7 = call float @llvm.nvvm.fma.rn.ftz.sat.f(float %0, float %1, float %6)
+ // CHECK-NEXT: ret float %7
+ // CHECK-NEXT: }
+ %f0 = nvvm.fma %a, %b, %c {rnd = #nvvm.fp_rnd_mode<rn>} : f32
+ %f1 = nvvm.fma %a, %b, %f0 {rnd = #nvvm.fp_rnd_mode<rn>, ftz = true} : f32
+ %f2 = nvvm.fma %a, %b, %f1 {rnd = #nvvm.fp_rnd_mode<rn>, sat = #nvvm.sat_mode<sat>} : f32
+ %f3 = nvvm.fma %a, %b, %f2 {rnd = #nvvm.fp_rnd_mode<rn>, sat = #nvvm.sat_mode<sat>, ftz = true} : f32
+ llvm.return %f3 : f32
+}
+
+llvm.func @fma_f32_rm(%a: f32, %b: f32, %c: f32) -> f32 {
+ // CHECK-LABEL: define float @fma_f32_rm(float %0, float %1, float %2) {
+ // CHECK-NEXT: %4 = call float @llvm.nvvm.fma.rm.f(float %0, float %1, float %2)
+ // CHECK-NEXT: %5 = call float @llvm.nvvm.fma.rm.ftz.f(float %0, float %1, float %4)
+ // CHECK-NEXT: %6 = call float @llvm.nvvm.fma.rm.sat.f(float %0, float %1, float %5)
+ // CHECK-NEXT: %7 = call float @llvm.nvvm.fma.rm.ftz.sat.f(float %0, float %1, float %6)
+ // CHECK-NEXT: ret float %7
+ // CHECK-NEXT: }
+ %f0 = nvvm.fma %a, %b, %c {rnd = #nvvm.fp_rnd_mode<rm>} : f32
+ %f1 = nvvm.fma %a, %b, %f0 {rnd = #nvvm.fp_rnd_mode<rm>, ftz = true} : f32
+ %f2 = nvvm.fma %a, %b, %f1 {rnd = #nvvm.fp_rnd_mode<rm>, sat = #nvvm.sat_mode<sat>} : f32
+ %f3 = nvvm.fma %a, %b, %f2 {rnd = #nvvm.fp_rnd_mode<rm>, sat = #nvvm.sat_mode<sat>, ftz = true} : f32
+ llvm.return %f3 : f32
+}
+
+llvm.func @fma_f32_rp(%a: f32, %b: f32, %c: f32) -> f32 {
+ // CHECK-LABEL: define float @fma_f32_rp(float %0, float %1, float %2) {
+ // CHECK-NEXT: %4 = call float @llvm.nvvm.fma.rp.f(float %0, float %1, float %2)
+ // CHECK-NEXT: %5 = call float @llvm.nvvm.fma.rp.ftz.f(float %0, float %1, float %4)
+ // CHECK-NEXT: %6 = call float @llvm.nvvm.fma.rp.sat.f(float %0, float %1, float %5)
+ // CHECK-NEXT: %7 = call float @llvm.nvvm.fma.rp.ftz.sat.f(float %0, float %1, float %6)
+ // CHECK-NEXT: ret float %7
+ // CHECK-NEXT: }
+ %f0 = nvvm.fma %a, %b, %c {rnd = #nvvm.fp_rnd_mode<rp>} : f32
+ %f1 = nvvm.fma %a, %b, %f0 {rnd = #nvvm.fp_rnd_mode<rp>, ftz = true} : f32
+ %f2 = nvvm.fma %a, %b, %f1 {rnd = #nvvm.fp_rnd_mode<rp>, sat = #nvvm.sat_mode<sat>} : f32
+ %f3 = nvvm.fma %a, %b, %f2 {rnd = #nvvm.fp_rnd_mode<rp>, sat = #nvvm.sat_mode<sat>, ftz = true} : f32
+ llvm.return %f3 : f32
+}
+
+llvm.func @fma_f32_rz(%a: f32, %b: f32, %c: f32) -> f32 {
+ // CHECK-LABEL: define float @fma_f32_rz(float %0, float %1, float %2) {
+ // CHECK-NEXT: %4 = call float @llvm.nvvm.fma.rz.f(float %0, float %1, float %2)
+ // CHECK-NEXT: %5 = call float @llvm.nvvm.fma.rz.ftz.f(float %0, float %1, float %4)
+ // CHECK-NEXT: %6 = call float @llvm.nvvm.fma.rz.sat.f(float %0, float %1, float %5)
+ // CHECK-NEXT: %7 = call float @llvm.nvvm.fma.rz.ftz.sat.f(float %0, float %1, float %6)
+ // CHECK-NEXT: ret float %7
+ // CHECK-NEXT: }
+ %f0 = nvvm.fma %a, %b, %c {rnd = #nvvm.fp_rnd_mode<rz>} : f32
+ %f1 = nvvm.fma %a, %b, %f0 {rnd = #nvvm.fp_rnd_mode<rz>, ftz = true} : f32
+ %f2 = nvvm.fma %a, %b, %f1 {rnd = #nvvm.fp_rnd_mode<rz>, sat = #nvvm.sat_mode<sat>} : f32
+ %f3 = nvvm.fma %a, %b, %f2 {rnd = #nvvm.fp_rnd_mode<rz>, sat = #nvvm.sat_mode<sat>, ftz = true} : f32
+ llvm.return %f3 : f32
+}
+
+llvm.func @fma_f64(%a: f64, %b: f64, %c: f64) -> f64 {
+ // CHECK-LABEL: define double @fma_f64(double %0, double %1, double %2) {
+ // CHECK-NEXT: %4 = call double @llvm.nvvm.fma.rn.d(double %0, double %1, double %2)
+ // CHECK-NEXT: %5 = call double @llvm.nvvm.fma.rm.d(double %0, double %1, double %4)
+ // CHECK-NEXT: %6 = call double @llvm.nvvm.fma.rp.d(double %0, double %1, double %5)
+ // CHECK-NEXT: %7 = call double @llvm.nvvm.fma.rz.d(double %0, double %1, double %6)
+ // CHECK-NEXT: ret double %7
+ // CHECK-NEXT: }
+ %f0 = nvvm.fma %a, %b, %c {rnd = #nvvm.fp_rnd_mode<rn>} : f64
+ %f1 = nvvm.fma %a, %b, %f0 {rnd = #nvvm.fp_rnd_mode<rm>} : f64
+ %f2 = nvvm.fma %a, %b, %f1 {rnd = #nvvm.fp_rnd_mode<rp>} : f64
+ %f3 = nvvm.fma %a, %b, %f2 {rnd = #nvvm.fp_rnd_mode<rz>} : f64
+ llvm.return %f3 : f64
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/fma/fma_invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/fma/fma_invalid.mlir
new file mode 100644
index 0000000000000..79b5c80abbe62
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/fma/fma_invalid.mlir
@@ -0,0 +1,89 @@
+// RUN: mlir-translate --mlir-to-llvmir --split-input-file --verify-diagnostics %s
+
+// -----
+
+llvm.func @fma_invalid_rnd_mode(%a : f16, %b : f16, %c : f16) -> f16 {
+ // expected-error at +1 {{rounding mode must be specified}}
+ %f1 = nvvm.fma %a, %b, %c {rnd = #nvvm.fp_rnd_mode<none>} : f16
+ llvm.return %f1 : f16
+}
+
+// -----
+
+llvm.func @fma_invalid_sat_mode(%a : f16, %b : f16, %c : f16) -> f16 {
+ // expected-error at +1 {{attribute 'sat' failed to satisfy constraint: Describes the saturation mode whose value is one of {none, sat}}}
+ %f1 = nvvm.fma %a, %b, %c {sat = #nvvm.sat_mode<satfinite>} : f16
+ llvm.return %f1 : f16
+}
+
+// -----
+
+llvm.func @fma_invalid_relu_sat(%a : f16, %b : f16, %c : f16) -> f16 {
+ // expected-error at +1 {{relu and saturation are not supported together}}
+ %f1 = nvvm.fma %a, %b, %c {relu = true, sat = #nvvm.sat_mode<sat>} : f16
+ llvm.return %f1 : f16
+}
+
+// -----
+
+llvm.func @fma_invalid_oob_sat(%a : f16, %b : f16, %c : f16) -> f16 {
+ // expected-error at +1 {{oob is not supported with saturation}}
+ %f1 = nvvm.fma %a, %b, %c {oob = true, sat = #nvvm.sat_mode<sat>} : f16
+ llvm.return %f1 : f16
+}
+
+// -----
+
+llvm.func @fma_invalid_oob_f64(%a : f64, %b : f64, %c : f64) -> f64 {
+ // expected-error at +1 {{oob is supported only for f16 and bf16 fused multiply-add operations}}
+ %f1 = nvvm.fma %a, %b, %c {oob = true} : f64
+ llvm.return %f1 : f64
+}
+
+// -----
+
+llvm.func @fma_invalid_relu_oob(%a : f16, %b : f16, %c : f16) -> f16 {
+ // expected-error at +1 {{relu and oob are only supported for f16 and bf16 fused multiply-add operations}}
+ %f1 = nvvm.fma %a, %b, %c {relu = true, oob = true} : f16
+ llvm.return %f1 : f16
+}
+
+// -----
+
+llvm.func @fma_invalid_ftz_sat_f64(%a : f64, %b : f64, %c : f64) -> f64 {
+ // expected-error at +1 {{FTZ and saturation are not supported for fused multiply-add operations involving f64 type}}
+ %f1 = nvvm.fma %a, %b, %c {ftz = true, sat = #nvvm.sat_mode<sat>} : f64
+ llvm.return %f1 : f64
+}
+
+// -----
+
+llvm.func @fma_invalid_v2f16_rnd_mode(%a : vector<2xf16>, %b : vector<2xf16>, %c : vector<2xf16>) -> vector<2xf16> {
+ // expected-error at +1 {{only RN rounding mode is supported for f16 and vector<2xf16> fused multiply-add operations}}
+ %f1 = nvvm.fma %a, %b, %c {rnd = #nvvm.fp_rnd_mode<rm>} : vector<2xf16>
+ llvm.return %f1 : vector<2xf16>
+}
+
+// -----
+
+llvm.func @fma_invalid_v2bf16_rnd_mode(%a : vector<2xbf16>, %b : vector<2xbf16>, %c : vector<2xbf16>) -> vector<2xbf16> {
+ // expected-error at +1 {{only RN rounding mode is supported for bf16 and vector<2xbf16> fused multiply-add operations}}
+ %f1 = nvvm.fma %a, %b, %c {rnd = #nvvm.fp_rnd_mode<rm>} : vector<2xbf16>
+ llvm.return %f1 : vector<2xbf16>
+}
+
+// -----
+
+llvm.func @fma_invalid_ftz_v2bf16(%a : vector<2xbf16>, %b : vector<2xbf16>, %c : vector<2xbf16>) -> vector<2xbf16> {
+ // expected-error at +1 {{FTZ and saturation are not supported for bf16 and vector<2xbf16> fused multiply-add operations}}
+ %f1 = nvvm.fma %a, %b, %c {ftz = true} : vector<2xbf16>
+ llvm.return %f1 : vector<2xbf16>
+}
+
+// -----
+
+llvm.func @fma_invalid_sat_v2bf16(%a : vector<2xbf16>, %b : vector<2xbf16>, %c : vector<2xbf16>) -> vector<2xbf16> {
+ // expected-error at +1 {{FTZ and saturation are not supported for bf16 and vector<2xbf16> fused multiply-add operations}}
+ %f1 = nvvm.fma %a, %b, %c {sat = #nvvm.sat_mode<sat>} : vector<2xbf16>
+ llvm.return %f1 : vector<2xbf16>
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/fma/fma_vector.mlir b/mlir/test/Target/LLVMIR/nvvm/fma/fma_vector.mlir
new file mode 100644
index 0000000000000..020bdcfc27705
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/fma/fma_vector.mlir
@@ -0,0 +1,294 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+llvm.func @fma_f16(%a: vector<2xf16>, %b: vector<2xf16>, %c: vector<2xf16>) -> vector<2xf16> {
+ // CHECK-LABEL: define <2 x half> @fma_f16(<2 x half> %0, <2 x half> %1, <2 x half> %2) {
+ // CHECK-NEXT: %4 = call <2 x half> @llvm.nvvm.fma.rn.f16x2(<2 x half> %0, <2 x half> %1, <2 x half> %2)
+ // CHECK-NEXT: %5 = call <2 x half> @llvm.nvvm.fma.rn.ftz.f16x2(<2 x half> %0, <2 x half> %1, <2 x half> %4)
+ // CHECK-NEXT: %6 = call <2 x half> @llvm.nvvm.fma.rn.sat.f16x2(<2 x half> %0, <2 x half> %1, <2 x half> %5)
+ // CHECK-NEXT: %7 = call <2 x half> @llvm.nvvm.fma.rn.ftz.sat.f16x2(<2 x half> %0, <2 x half> %1, <2 x half> %6)
+ // CHECK-NEXT: %8 = call <2 x half> @llvm.nvvm.fma.rn.relu.f16x2(<2 x half> %0, <2 x half> %1, <2 x half> %7)
+ // CHECK-NEXT: %9 = call <2 x half> @llvm.nvvm.fma.rn.ftz.relu.f16x2(<2 x half> %0, <2 x half> %1, <2 x half> %8)
+ // CHECK-NEXT: %10 = call <2 x half> @llvm.nvvm.fma.rn.oob.v2f16(<2 x half> %0, <2 x half> %1, <2 x half> %9)
+ // CHECK-NEXT: %11 = call <2 x half> @llvm.nvvm.fma.rn.oob.relu.v2f16(<2 x half> %0, <2 x half> %1, <2 x half> %10)
+ // CHECK-NEXT: ret <2 x half> %11
+ // CHECK-NEXT: }
+ %f0 = nvvm.fma %a, %b, %c {rnd = #nvvm.fp_rnd_mode<rn>} : vector<2xf16>
+ %f1 = nvvm.fma %a, %b, %f0 {rnd = #nvvm.fp_rnd_mode<rn>, ftz = true} : vector<2xf16>
+ %f2 = nvvm.fma %a, %b, %f1 {rnd = #nvvm.fp_rnd_mode<rn>, sat = #nvvm.sat_mode<sat>} : vector<2xf16>
+ %f3 = nvvm.fma %a, %b, %f2 {rnd = #nvvm.fp_rnd_mode<rn>, sat = #nvvm.sat_mode<sat>, ftz = true} : vector<2xf16>
+ %f4 = nvvm.fma %a, %b, %f3 {rnd = #nvvm.fp_rnd_mode<rn>, relu = true} : vector<2xf16>
+ %f5 = nvvm.fma %a, %b, %f4 {rnd = #nvvm.fp_rnd_mode<rn>, relu = true, ftz = true} : vector<2xf16>
+ %f6 = nvvm.fma %a, %b, %f5 {rnd = #nvvm.fp_rnd_mode<rn>, oob = true} : vector<2xf16>
+ %f7 = nvvm.fma %a, %b, %f6 {rnd = #nvvm.fp_rnd_mode<rn>, oob = true, relu = true} : vector<2xf16>
+ llvm.return %f7 : vector<2xf16>
+}
+
+llvm.func @fma_bf16(%a: vector<2xbf16>, %b: vector<2xbf16>, %c: vector<2xbf16>) -> vector<2xbf16> {
+ // CHECK-LABEL: define <2 x bfloat> @fma_bf16(<2 x bfloat> %0, <2 x bfloat> %1, <2 x bfloat> %2) {
+ // CHECK-NEXT: %4 = call <2 x bfloat> @llvm.nvvm.fma.rn.bf16x2(<2 x bfloat> %0, <2 x bfloat> %1, <2 x bfloat> %2)
+ // CHECK-NEXT: %5 = call <2 x bfloat> @llvm.nvvm.fma.rn.relu.bf16x2(<2 x bfloat> %0, <2 x bfloat> %1, <2 x bfloat> %4)
+ // CHECK-NEXT: %6 = call <2 x bfloat> @llvm.nvvm.fma.rn.oob.v2bf16(<2 x bfloat> %0, <2 x bfloat> %1, <2 x bfloat> %5)
+ // CHECK-NEXT: %7 = call <2 x bfloat> @llvm.nvvm.fma.rn.oob.relu.v2bf16(<2 x bfloat> %0, <2 x bfloat> %1, <2 x bfloat> %6)
+ // CHECK-NEXT: ret <2 x bfloat> %7
+ // CHECK-NEXT: }
+ %f0 = nvvm.fma %a, %b, %c {rnd = #nvvm.fp_rnd_mode<rn>} : vector<2xbf16>
+ %f1 = nvvm.fma %a, %b, %f0 {rnd = #nvvm.fp_rnd_mode<rn>, relu = true} : vector<2xbf16>
+ %f2 = nvvm.fma %a, %b, %f1 {rnd = #nvvm.fp_rnd_mode<rn>, oob = true} : vector<2xbf16>
+ %f3 = nvvm.fma %a, %b, %f2 {rnd = #nvvm.fp_rnd_mode<rn>, oob = true, relu = true} : vector<2xbf16>
+ llvm.return %f3 : vector<2xbf16>
+}
+
+llvm.func @fma_f32_rn(%a: vector<2xf32>, %b: vector<2xf32>, %c: vector<2xf32>) -> vector<2xf32> {
+ // CHECK-LABEL: define <2 x float> @fma_f32_rn(<2 x float> %0, <2 x float> %1, <2 x float> %2) {
+ // CHECK-NEXT: %4 = extractelement <2 x float> %0, i32 0
+ // CHECK-NEXT: %5 = extractelement <2 x float> %1, i32 0
+ // CHECK-NEXT: %6 = extractelement <2 x float> %2, i32 0
+ // CHECK-NEXT: %7 = call float @llvm.nvvm.fma.rn.f(float %4, float %5, float %6)
+ // CHECK-NEXT: %8 = insertelement <2 x float> poison, float %7, i32 0
+ // CHECK-NEXT: %9 = extractelement <2 x float> %0, i32 1
+ // CHECK-NEXT: %10 = extractelement <2 x float> %1, i32 1
+ // CHECK-NEXT: %11 = extractelement <2 x float> %2, i32 1
+ // CHECK-NEXT: %12 = call float @llvm.nvvm.fma.rn.f(float %9, float %10, float %11)
+ // CHECK-NEXT: %13 = insertelement <2 x float> %8, float %12, i32 1
+ // CHECK-NEXT: %14 = extractelement <2 x float> %0, i32 0
+ // CHECK-NEXT: %15 = extractelement <2 x float> %1, i32 0
+ // CHECK-NEXT: %16 = extractelement <2 x float> %13, i32 0
+ // CHECK-NEXT: %17 = call float @llvm.nvvm.fma.rn.ftz.f(float %14, float %15, float %16)
+ // CHECK-NEXT: %18 = insertelement <2 x float> poison, float %17, i32 0
+ // CHECK-NEXT: %19 = extractelement <2 x float> %0, i32 1
+ // CHECK-NEXT: %20 = extractelement <2 x float> %1, i32 1
+ // CHECK-NEXT: %21 = extractelement <2 x float> %13, i32 1
+ // CHECK-NEXT: %22 = call float @llvm.nvvm.fma.rn.ftz.f(float %19, float %20, float %21)
+ // CHECK-NEXT: %23 = insertelement <2 x float> %18, float %22, i32 1
+ // CHECK-NEXT: %24 = extractelement <2 x float> %0, i32 0
+ // CHECK-NEXT: %25 = extractelement <2 x float> %1, i32 0
+ // CHECK-NEXT: %26 = extractelement <2 x float> %23, i32 0
+ // CHECK-NEXT: %27 = call float @llvm.nvvm.fma.rn.sat.f(float %24, float %25, float %26)
+ // CHECK-NEXT: %28 = insertelement <2 x float> poison, float %27, i32 0
+ // CHECK-NEXT: %29 = extractelement <2 x float> %0, i32 1
+ // CHECK-NEXT: %30 = extractelement <2 x float> %1, i32 1
+ // CHECK-NEXT: %31 = extractelement <2 x float> %23, i32 1
+ // CHECK-NEXT: %32 = call float @llvm.nvvm.fma.rn.sat.f(float %29, float %30, float %31)
+ // CHECK-NEXT: %33 = insertelement <2 x float> %28, float %32, i32 1
+ // CHECK-NEXT: %34 = extractelement <2 x float> %0, i32 0
+ // CHECK-NEXT: %35 = extractelement <2 x float> %1, i32 0
+ // CHECK-NEXT: %36 = extractelement <2 x float> %33, i32 0
+ // CHECK-NEXT: %37 = call float @llvm.nvvm.fma.rn.ftz.sat.f(float %34, float %35, float %36)
+ // CHECK-NEXT: %38 = insertelement <2 x float> poison, float %37, i32 0
+ // CHECK-NEXT: %39 = extractelement <2 x float> %0, i32 1
+ // CHECK-NEXT: %40 = extractelement <2 x float> %1, i32 1
+ // CHECK-NEXT: %41 = extractelement <2 x float> %33, i32 1
+ // CHECK-NEXT: %42 = call float @llvm.nvvm.fma.rn.ftz.sat.f(float %39, float %40, float %41)
+ // CHECK-NEXT: %43 = insertelement <2 x float> %38, float %42, i32 1
+ // CHECK-NEXT: ret <2 x float> %43
+ // CHECK-NEXT: }
+ %f0 = nvvm.fma %a, %b, %c {rnd = #nvvm.fp_rnd_mode<rn>} : vector<2xf32>
+ %f1 = nvvm.fma %a, %b, %f0 {rnd = #nvvm.fp_rnd_mode<rn>, ftz = true} : vector<2xf32>
+ %f2 = nvvm.fma %a, %b, %f1 {rnd = #nvvm.fp_rnd_mode<rn>, sat = #nvvm.sat_mode<sat>} : vector<2xf32>
+ %f3 = nvvm.fma %a, %b, %f2 {rnd = #nvvm.fp_rnd_mode<rn>, sat = #nvvm.sat_mode<sat>, ftz = true} : vector<2xf32>
+ llvm.return %f3 : vector<2xf32>
+}
+
+llvm.func @fma_f32_rm(%a: vector<2xf32>, %b: vector<2xf32>, %c: vector<2xf32>) -> vector<2xf32> {
+ // CHECK-LABEL: define <2 x float> @fma_f32_rm(<2 x float> %0, <2 x float> %1, <2 x float> %2) {
+ // CHECK-NEXT: %4 = extractelement <2 x float> %0, i32 0
+ // CHECK-NEXT: %5 = extractelement <2 x float> %1, i32 0
+ // CHECK-NEXT: %6 = extractelement <2 x float> %2, i32 0
+ // CHECK-NEXT: %7 = call float @llvm.nvvm.fma.rm.f(float %4, float %5, float %6)
+ // CHECK-NEXT: %8 = insertelement <2 x float> poison, float %7, i32 0
+ // CHECK-NEXT: %9 = extractelement <2 x float> %0, i32 1
+ // CHECK-NEXT: %10 = extractelement <2 x float> %1, i32 1
+ // CHECK-NEXT: %11 = extractelement <2 x float> %2, i32 1
+ // CHECK-NEXT: %12 = call float @llvm.nvvm.fma.rm.f(float %9, float %10, float %11)
+ // CHECK-NEXT: %13 = insertelement <2 x float> %8, float %12, i32 1
+ // CHECK-NEXT: %14 = extractelement <2 x float> %0, i32 0
+ // CHECK-NEXT: %15 = extractelement <2 x float> %1, i32 0
+ // CHECK-NEXT: %16 = extractelement <2 x float> %13, i32 0
+ // CHECK-NEXT: %17 = call float @llvm.nvvm.fma.rm.ftz.f(float %14, float %15, float %16)
+ // CHECK-NEXT: %18 = insertelement <2 x float> poison, float %17, i32 0
+ // CHECK-NEXT: %19 = extractelement <2 x float> %0, i32 1
+ // CHECK-NEXT: %20 = extractelement <2 x float> %1, i32 1
+ // CHECK-NEXT: %21 = extractelement <2 x float> %13, i32 1
+ // CHECK-NEXT: %22 = call float @llvm.nvvm.fma.rm.ftz.f(float %19, float %20, float %21)
+ // CHECK-NEXT: %23 = insertelement <2 x float> %18, float %22, i32 1
+ // CHECK-NEXT: %24 = extractelement <2 x float> %0, i32 0
+ // CHECK-NEXT: %25 = extractelement <2 x float> %1, i32 0
+ // CHECK-NEXT: %26 = extractelement <2 x float> %23, i32 0
+ // CHECK-NEXT: %27 = call float @llvm.nvvm.fma.rm.sat.f(float %24, float %25, float %26)
+ // CHECK-NEXT: %28 = insertelement <2 x float> poison, float %27, i32 0
+ // CHECK-NEXT: %29 = extractelement <2 x float> %0, i32 1
+ // CHECK-NEXT: %30 = extractelement <2 x float> %1, i32 1
+ // CHECK-NEXT: %31 = extractelement <2 x float> %23, i32 1
+ // CHECK-NEXT: %32 = call float @llvm.nvvm.fma.rm.sat.f(float %29, float %30, float %31)
+ // CHECK-NEXT: %33 = insertelement <2 x float> %28, float %32, i32 1
+ // CHECK-NEXT: %34 = extractelement <2 x float> %0, i32 0
+ // CHECK-NEXT: %35 = extractelement <2 x float> %1, i32 0
+ // CHECK-NEXT: %36 = extractelement <2 x float> %33, i32 0
+ // CHECK-NEXT: %37 = call float @llvm.nvvm.fma.rm.ftz.sat.f(float %34, float %35, float %36)
+ // CHECK-NEXT: %38 = insertelement <2 x float> poison, float %37, i32 0
+ // CHECK-NEXT: %39 = extractelement <2 x float> %0, i32 1
+ // CHECK-NEXT: %40 = extractelement <2 x float> %1, i32 1
+ // CHECK-NEXT: %41 = extractelement <2 x float> %33, i32 1
+ // CHECK-NEXT: %42 = call float @llvm.nvvm.fma.rm.ftz.sat.f(float %39, float %40, float %41)
+ // CHECK-NEXT: %43 = insertelement <2 x float> %38, float %42, i32 1
+ // CHECK-NEXT: ret <2 x float> %43
+ // CHECK-NEXT: }
+ %f0 = nvvm.fma %a, %b, %c {rnd = #nvvm.fp_rnd_mode<rm>} : vector<2xf32>
+ %f1 = nvvm.fma %a, %b, %f0 {rnd = #nvvm.fp_rnd_mode<rm>, ftz = true} : vector<2xf32>
+ %f2 = nvvm.fma %a, %b, %f1 {rnd = #nvvm.fp_rnd_mode<rm>, sat = #nvvm.sat_mode<sat>} : vector<2xf32>
+ %f3 = nvvm.fma %a, %b, %f2 {rnd = #nvvm.fp_rnd_mode<rm>, sat = #nvvm.sat_mode<sat>, ftz = true} : vector<2xf32>
+ llvm.return %f3 : vector<2xf32>
+}
+
+llvm.func @fma_f32_rp(%a: vector<2xf32>, %b: vector<2xf32>, %c: vector<2xf32>) -> vector<2xf32> {
+ // CHECK-LABEL: define <2 x float> @fma_f32_rp(<2 x float> %0, <2 x float> %1, <2 x float> %2) {
+ // CHECK-NEXT: %4 = extractelement <2 x float> %0, i32 0
+ // CHECK-NEXT: %5 = extractelement <2 x float> %1, i32 0
+ // CHECK-NEXT: %6 = extractelement <2 x float> %2, i32 0
+ // CHECK-NEXT: %7 = call float @llvm.nvvm.fma.rp.f(float %4, float %5, float %6)
+ // CHECK-NEXT: %8 = insertelement <2 x float> poison, float %7, i32 0
+ // CHECK-NEXT: %9 = extractelement <2 x float> %0, i32 1
+ // CHECK-NEXT: %10 = extractelement <2 x float> %1, i32 1
+ // CHECK-NEXT: %11 = extractelement <2 x float> %2, i32 1
+ // CHECK-NEXT: %12 = call float @llvm.nvvm.fma.rp.f(float %9, float %10, float %11)
+ // CHECK-NEXT: %13 = insertelement <2 x float> %8, float %12, i32 1
+ // CHECK-NEXT: %14 = extractelement <2 x float> %0, i32 0
+ // CHECK-NEXT: %15 = extractelement <2 x float> %1, i32 0
+ // CHECK-NEXT: %16 = extractelement <2 x float> %13, i32 0
+ // CHECK-NEXT: %17 = call float @llvm.nvvm.fma.rp.ftz.f(float %14, float %15, float %16)
+ // CHECK-NEXT: %18 = insertelement <2 x float> poison, float %17, i32 0
+ // CHECK-NEXT: %19 = extractelement <2 x float> %0, i32 1
+ // CHECK-NEXT: %20 = extractelement <2 x float> %1, i32 1
+ // CHECK-NEXT: %21 = extractelement <2 x float> %13, i32 1
+ // CHECK-NEXT: %22 = call float @llvm.nvvm.fma.rp.ftz.f(float %19, float %20, float %21)
+ // CHECK-NEXT: %23 = insertelement <2 x float> %18, float %22, i32 1
+ // CHECK-NEXT: %24 = extractelement <2 x float> %0, i32 0
+ // CHECK-NEXT: %25 = extractelement <2 x float> %1, i32 0
+ // CHECK-NEXT: %26 = extractelement <2 x float> %23, i32 0
+ // CHECK-NEXT: %27 = call float @llvm.nvvm.fma.rp.sat.f(float %24, float %25, float %26)
+ // CHECK-NEXT: %28 = insertelement <2 x float> poison, float %27, i32 0
+ // CHECK-NEXT: %29 = extractelement <2 x float> %0, i32 1
+ // CHECK-NEXT: %30 = extractelement <2 x float> %1, i32 1
+ // CHECK-NEXT: %31 = extractelement <2 x float> %23, i32 1
+ // CHECK-NEXT: %32 = call float @llvm.nvvm.fma.rp.sat.f(float %29, float %30, float %31)
+ // CHECK-NEXT: %33 = insertelement <2 x float> %28, float %32, i32 1
+ // CHECK-NEXT: %34 = extractelement <2 x float> %0, i32 0
+ // CHECK-NEXT: %35 = extractelement <2 x float> %1, i32 0
+ // CHECK-NEXT: %36 = extractelement <2 x float> %33, i32 0
+ // CHECK-NEXT: %37 = call float @llvm.nvvm.fma.rp.ftz.sat.f(float %34, float %35, float %36)
+ // CHECK-NEXT: %38 = insertelement <2 x float> poison, float %37, i32 0
+ // CHECK-NEXT: %39 = extractelement <2 x float> %0, i32 1
+ // CHECK-NEXT: %40 = extractelement <2 x float> %1, i32 1
+ // CHECK-NEXT: %41 = extractelement <2 x float> %33, i32 1
+ // CHECK-NEXT: %42 = call float @llvm.nvvm.fma.rp.ftz.sat.f(float %39, float %40, float %41)
+ // CHECK-NEXT: %43 = insertelement <2 x float> %38, float %42, i32 1
+ // CHECK-NEXT: ret <2 x float> %43
+ // CHECK-NEXT: }
+ %f0 = nvvm.fma %a, %b, %c {rnd = #nvvm.fp_rnd_mode<rp>} : vector<2xf32>
+ %f1 = nvvm.fma %a, %b, %f0 {rnd = #nvvm.fp_rnd_mode<rp>, ftz = true} : vector<2xf32>
+ %f2 = nvvm.fma %a, %b, %f1 {rnd = #nvvm.fp_rnd_mode<rp>, sat = #nvvm.sat_mode<sat>} : vector<2xf32>
+ %f3 = nvvm.fma %a, %b, %f2 {rnd = #nvvm.fp_rnd_mode<rp>, sat = #nvvm.sat_mode<sat>, ftz = true} : vector<2xf32>
+ llvm.return %f3 : vector<2xf32>
+}
+
+llvm.func @fma_f32_rz(%a: vector<2xf32>, %b: vector<2xf32>, %c: vector<2xf32>) -> vector<2xf32> {
+ // CHECK-LABEL: define <2 x float> @fma_f32_rz(<2 x float> %0, <2 x float> %1, <2 x float> %2) {
+ // CHECK-NEXT: %4 = extractelement <2 x float> %0, i32 0
+ // CHECK-NEXT: %5 = extractelement <2 x float> %1, i32 0
+ // CHECK-NEXT: %6 = extractelement <2 x float> %2, i32 0
+ // CHECK-NEXT: %7 = call float @llvm.nvvm.fma.rz.f(float %4, float %5, float %6)
+ // CHECK-NEXT: %8 = insertelement <2 x float> poison, float %7, i32 0
+ // CHECK-NEXT: %9 = extractelement <2 x float> %0, i32 1
+ // CHECK-NEXT: %10 = extractelement <2 x float> %1, i32 1
+ // CHECK-NEXT: %11 = extractelement <2 x float> %2, i32 1
+ // CHECK-NEXT: %12 = call float @llvm.nvvm.fma.rz.f(float %9, float %10, float %11)
+ // CHECK-NEXT: %13 = insertelement <2 x float> %8, float %12, i32 1
+ // CHECK-NEXT: %14 = extractelement <2 x float> %0, i32 0
+ // CHECK-NEXT: %15 = extractelement <2 x float> %1, i32 0
+ // CHECK-NEXT: %16 = extractelement <2 x float> %13, i32 0
+ // CHECK-NEXT: %17 = call float @llvm.nvvm.fma.rz.ftz.f(float %14, float %15, float %16)
+ // CHECK-NEXT: %18 = insertelement <2 x float> poison, float %17, i32 0
+ // CHECK-NEXT: %19 = extractelement <2 x float> %0, i32 1
+ // CHECK-NEXT: %20 = extractelement <2 x float> %1, i32 1
+ // CHECK-NEXT: %21 = extractelement <2 x float> %13, i32 1
+ // CHECK-NEXT: %22 = call float @llvm.nvvm.fma.rz.ftz.f(float %19, float %20, float %21)
+ // CHECK-NEXT: %23 = insertelement <2 x float> %18, float %22, i32 1
+ // CHECK-NEXT: %24 = extractelement <2 x float> %0, i32 0
+ // CHECK-NEXT: %25 = extractelement <2 x float> %1, i32 0
+ // CHECK-NEXT: %26 = extractelement <2 x float> %23, i32 0
+ // CHECK-NEXT: %27 = call float @llvm.nvvm.fma.rz.sat.f(float %24, float %25, float %26)
+ // CHECK-NEXT: %28 = insertelement <2 x float> poison, float %27, i32 0
+ // CHECK-NEXT: %29 = extractelement <2 x float> %0, i32 1
+ // CHECK-NEXT: %30 = extractelement <2 x float> %1, i32 1
+ // CHECK-NEXT: %31 = extractelement <2 x float> %23, i32 1
+ // CHECK-NEXT: %32 = call float @llvm.nvvm.fma.rz.sat.f(float %29, float %30, float %31)
+ // CHECK-NEXT: %33 = insertelement <2 x float> %28, float %32, i32 1
+ // CHECK-NEXT: %34 = extractelement <2 x float> %0, i32 0
+ // CHECK-NEXT: %35 = extractelement <2 x float> %1, i32 0
+ // CHECK-NEXT: %36 = extractelement <2 x float> %33, i32 0
+ // CHECK-NEXT: %37 = call float @llvm.nvvm.fma.rz.ftz.sat.f(float %34, float %35, float %36)
+ // CHECK-NEXT: %38 = insertelement <2 x float> poison, float %37, i32 0
+ // CHECK-NEXT: %39 = extractelement <2 x float> %0, i32 1
+ // CHECK-NEXT: %40 = extractelement <2 x float> %1, i32 1
+ // CHECK-NEXT: %41 = extractelement <2 x float> %33, i32 1
+ // CHECK-NEXT: %42 = call float @llvm.nvvm.fma.rz.ftz.sat.f(float %39, float %40, float %41)
+ // CHECK-NEXT: %43 = insertelement <2 x float> %38, float %42, i32 1
+ // CHECK-NEXT: ret <2 x float> %43
+ // CHECK-NEXT: }
+ %f0 = nvvm.fma %a, %b, %c {rnd = #nvvm.fp_rnd_mode<rz>} : vector<2xf32>
+ %f1 = nvvm.fma %a, %b, %f0 {rnd = #nvvm.fp_rnd_mode<rz>, ftz = true} : vector<2xf32>
+ %f2 = nvvm.fma %a, %b, %f1 {rnd = #nvvm.fp_rnd_mode<rz>, sat = #nvvm.sat_mode<sat>} : vector<2xf32>
+ %f3 = nvvm.fma %a, %b, %f2 {rnd = #nvvm.fp_rnd_mode<rz>, sat = #nvvm.sat_mode<sat>, ftz = true} : vector<2xf32>
+ llvm.return %f3 : vector<2xf32>
+}
+
+llvm.func @fma_f64(%a: vector<2xf64>, %b: vector<2xf64>, %c: vector<2xf64>) -> vector<2xf64> {
+ // CHECK-LABEL: define <2 x double> @fma_f64(<2 x double> %0, <2 x double> %1, <2 x double> %2) {
+ // CHECK-NEXT: %4 = extractelement <2 x double> %0, i32 0
+ // CHECK-NEXT: %5 = extractelement <2 x double> %1, i32 0
+ // CHECK-NEXT: %6 = extractelement <2 x double> %2, i32 0
+ // CHECK-NEXT: %7 = call double @llvm.nvvm.fma.rn.d(double %4, double %5, double %6)
+ // CHECK-NEXT: %8 = insertelement <2 x double> poison, double %7, i32 0
+ // CHECK-NEXT: %9 = extractelement <2 x double> %0, i32 1
+ // CHECK-NEXT: %10 = extractelement <2 x double> %1, i32 1
+ // CHECK-NEXT: %11 = extractelement <2 x double> %2, i32 1
+ // CHECK-NEXT: %12 = call double @llvm.nvvm.fma.rn.d(double %9, double %10, double %11)
+ // CHECK-NEXT: %13 = insertelement <2 x double> %8, double %12, i32 1
+ // CHECK-NEXT: %14 = extractelement <2 x double> %0, i32 0
+ // CHECK-NEXT: %15 = extractelement <2 x double> %1, i32 0
+ // CHECK-NEXT: %16 = extractelement <2 x double> %13, i32 0
+ // CHECK-NEXT: %17 = call double @llvm.nvvm.fma.rm.d(double %14, double %15, double %16)
+ // CHECK-NEXT: %18 = insertelement <2 x double> poison, double %17, i32 0
+ // CHECK-NEXT: %19 = extractelement <2 x double> %0, i32 1
+ // CHECK-NEXT: %20 = extractelement <2 x double> %1, i32 1
+ // CHECK-NEXT: %21 = extractelement <2 x double> %13, i32 1
+ // CHECK-NEXT: %22 = call double @llvm.nvvm.fma.rm.d(double %19, double %20, double %21)
+ // CHECK-NEXT: %23 = insertelement <2 x double> %18, double %22, i32 1
+ // CHECK-NEXT: %24 = extractelement <2 x double> %0, i32 0
+ // CHECK-NEXT: %25 = extractelement <2 x double> %1, i32 0
+ // CHECK-NEXT: %26 = extractelement <2 x double> %23, i32 0
+ // CHECK-NEXT: %27 = call double @llvm.nvvm.fma.rp.d(double %24, double %25, double %26)
+ // CHECK-NEXT: %28 = insertelement <2 x double> poison, double %27, i32 0
+ // CHECK-NEXT: %29 = extractelement <2 x double> %0, i32 1
+ // CHECK-NEXT: %30 = extractelement <2 x double> %1, i32 1
+ // CHECK-NEXT: %31 = extractelement <2 x double> %23, i32 1
+ // CHECK-NEXT: %32 = call double @llvm.nvvm.fma.rp.d(double %29, double %30, double %31)
+ // CHECK-NEXT: %33 = insertelement <2 x double> %28, double %32, i32 1
+ // CHECK-NEXT: %34 = extractelement <2 x double> %0, i32 0
+ // CHECK-NEXT: %35 = extractelement <2 x double> %1, i32 0
+ // CHECK-NEXT: %36 = extractelement <2 x double> %33, i32 0
+ // CHECK-NEXT: %37 = call double @llvm.nvvm.fma.rz.d(double %34, double %35, double %36)
+ // CHECK-NEXT: %38 = insertelement <2 x double> poison, double %37, i32 0
+ // CHECK-NEXT: %39 = extractelement <2 x double> %0, i32 1
+ // CHECK-NEXT: %40 = extractelement <2 x double> %1, i32 1
+ // CHECK-NEXT: %41 = extractelement <2 x double> %33, i32 1
+ // CHECK-NEXT: %42 = call double @llvm.nvvm.fma.rz.d(double %39, double %40, double %41)
+ // CHECK-NEXT: %43 = insertelement <2 x double> %38, double %42, i32 1
+ // CHECK-NEXT: ret <2 x double> %43
+ // CHECK-NEXT: }
+ %f0 = nvvm.fma %a, %b, %c {rnd = #nvvm.fp_rnd_mode<rn>} : vector<2xf64>
+ %f1 = nvvm.fma %a, %b, %f0 {rnd = #nvvm.fp_rnd_mode<rm>} : vector<2xf64>
+ %f2 = nvvm.fma %a, %b, %f1 {rnd = #nvvm.fp_rnd_mode<rp>} : vector<2xf64>
+ %f3 = nvvm.fma %a, %b, %f2 {rnd = #nvvm.fp_rnd_mode<rz>} : vector<2xf64>
+ llvm.return %f3 : vector<2xf64>
+}
>From 929079f33240602a595b231360f25366a767af35 Mon Sep 17 00:00:00 2001
From: Srinivasa Ravi <srinivasar at nvidia.com>
Date: Thu, 5 Mar 2026 12:32:27 +0000
Subject: [PATCH 2/2] fix test
---
.../test/Target/LLVMIR/nvvm/fma/fma_invalid.mlir | 16 ++++++++--------
1 file changed, 8 insertions(+), 8 deletions(-)
diff --git a/mlir/test/Target/LLVMIR/nvvm/fma/fma_invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/fma/fma_invalid.mlir
index 79b5c80abbe62..1d68ddef6c1ed 100644
--- a/mlir/test/Target/LLVMIR/nvvm/fma/fma_invalid.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/fma/fma_invalid.mlir
@@ -12,7 +12,7 @@ llvm.func @fma_invalid_rnd_mode(%a : f16, %b : f16, %c : f16) -> f16 {
llvm.func @fma_invalid_sat_mode(%a : f16, %b : f16, %c : f16) -> f16 {
// expected-error at +1 {{attribute 'sat' failed to satisfy constraint: Describes the saturation mode whose value is one of {none, sat}}}
- %f1 = nvvm.fma %a, %b, %c {sat = #nvvm.sat_mode<satfinite>} : f16
+ %f1 = nvvm.fma %a, %b, %c {sat = #nvvm.sat_mode<satfinite>, rnd = #nvvm.fp_rnd_mode<rn>} : f16
llvm.return %f1 : f16
}
@@ -20,7 +20,7 @@ llvm.func @fma_invalid_sat_mode(%a : f16, %b : f16, %c : f16) -> f16 {
llvm.func @fma_invalid_relu_sat(%a : f16, %b : f16, %c : f16) -> f16 {
// expected-error at +1 {{relu and saturation are not supported together}}
- %f1 = nvvm.fma %a, %b, %c {relu = true, sat = #nvvm.sat_mode<sat>} : f16
+ %f1 = nvvm.fma %a, %b, %c {relu = true, sat = #nvvm.sat_mode<sat>, rnd = #nvvm.fp_rnd_mode<rn>} : f16
llvm.return %f1 : f16
}
@@ -28,7 +28,7 @@ llvm.func @fma_invalid_relu_sat(%a : f16, %b : f16, %c : f16) -> f16 {
llvm.func @fma_invalid_oob_sat(%a : f16, %b : f16, %c : f16) -> f16 {
// expected-error at +1 {{oob is not supported with saturation}}
- %f1 = nvvm.fma %a, %b, %c {oob = true, sat = #nvvm.sat_mode<sat>} : f16
+ %f1 = nvvm.fma %a, %b, %c {oob = true, sat = #nvvm.sat_mode<sat>, rnd = #nvvm.fp_rnd_mode<rn>} : f16
llvm.return %f1 : f16
}
@@ -36,7 +36,7 @@ llvm.func @fma_invalid_oob_sat(%a : f16, %b : f16, %c : f16) -> f16 {
llvm.func @fma_invalid_oob_f64(%a : f64, %b : f64, %c : f64) -> f64 {
// expected-error at +1 {{oob is supported only for f16 and bf16 fused multiply-add operations}}
- %f1 = nvvm.fma %a, %b, %c {oob = true} : f64
+ %f1 = nvvm.fma %a, %b, %c {oob = true, rnd = #nvvm.fp_rnd_mode<rn>} : f64
llvm.return %f1 : f64
}
@@ -44,7 +44,7 @@ llvm.func @fma_invalid_oob_f64(%a : f64, %b : f64, %c : f64) -> f64 {
llvm.func @fma_invalid_relu_oob(%a : f16, %b : f16, %c : f16) -> f16 {
// expected-error at +1 {{relu and oob are only supported for f16 and bf16 fused multiply-add operations}}
- %f1 = nvvm.fma %a, %b, %c {relu = true, oob = true} : f16
+ %f1 = nvvm.fma %a, %b, %c {relu = true, oob = true, rnd = #nvvm.fp_rnd_mode<rn>} : f16
llvm.return %f1 : f16
}
@@ -52,7 +52,7 @@ llvm.func @fma_invalid_relu_oob(%a : f16, %b : f16, %c : f16) -> f16 {
llvm.func @fma_invalid_ftz_sat_f64(%a : f64, %b : f64, %c : f64) -> f64 {
// expected-error at +1 {{FTZ and saturation are not supported for fused multiply-add operations involving f64 type}}
- %f1 = nvvm.fma %a, %b, %c {ftz = true, sat = #nvvm.sat_mode<sat>} : f64
+ %f1 = nvvm.fma %a, %b, %c {ftz = true, sat = #nvvm.sat_mode<sat>, rnd = #nvvm.fp_rnd_mode<rn>} : f64
llvm.return %f1 : f64
}
@@ -76,7 +76,7 @@ llvm.func @fma_invalid_v2bf16_rnd_mode(%a : vector<2xbf16>, %b : vector<2xbf16>,
llvm.func @fma_invalid_ftz_v2bf16(%a : vector<2xbf16>, %b : vector<2xbf16>, %c : vector<2xbf16>) -> vector<2xbf16> {
// expected-error at +1 {{FTZ and saturation are not supported for bf16 and vector<2xbf16> fused multiply-add operations}}
- %f1 = nvvm.fma %a, %b, %c {ftz = true} : vector<2xbf16>
+ %f1 = nvvm.fma %a, %b, %c {ftz = true, rnd = #nvvm.fp_rnd_mode<rn>} : vector<2xbf16>
llvm.return %f1 : vector<2xbf16>
}
@@ -84,6 +84,6 @@ llvm.func @fma_invalid_ftz_v2bf16(%a : vector<2xbf16>, %b : vector<2xbf16>, %c :
llvm.func @fma_invalid_sat_v2bf16(%a : vector<2xbf16>, %b : vector<2xbf16>, %c : vector<2xbf16>) -> vector<2xbf16> {
// expected-error at +1 {{FTZ and saturation are not supported for bf16 and vector<2xbf16> fused multiply-add operations}}
- %f1 = nvvm.fma %a, %b, %c {sat = #nvvm.sat_mode<sat>} : vector<2xbf16>
+ %f1 = nvvm.fma %a, %b, %c {sat = #nvvm.sat_mode<sat>, rnd = #nvvm.fp_rnd_mode<rn>} : vector<2xbf16>
llvm.return %f1 : vector<2xbf16>
}
More information about the Mlir-commits
mailing list