[llvm] [SelectionDAG] Add `STRICT_BF16_TO_FP` and `STRICT_FP_TO_BF16` (PR #80056)
Shilei Tian via llvm-commits
llvm-commits at lists.llvm.org
Fri Feb 9 16:28:16 PST 2024
https://github.com/shiltian updated https://github.com/llvm/llvm-project/pull/80056
>From e71b92d4e8227cde08fc10fd6f3ff183ed638d90 Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Fri, 9 Feb 2024 19:25:41 -0500
Subject: [PATCH] [SelectionDAG] Add `STRICT_BF16_TO_FP` and
`STRICT_FP_TO_BF16`
This patch adds the support for `STRICT_BF16_TO_FP` and `STRICT_FP_TO_BF16`.
Fix #78540.
---
llvm/include/llvm/CodeGen/ISDOpcodes.h | 2 ++
llvm/include/llvm/CodeGen/SelectionDAGNodes.h | 2 ++
llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp | 32 +++++++++++++------
.../SelectionDAG/LegalizeFloatTypes.cpp | 25 +++++++++------
.../SelectionDAG/LegalizeIntegerTypes.cpp | 1 +
.../SelectionDAG/SelectionDAGDumper.cpp | 2 ++
6 files changed, 45 insertions(+), 19 deletions(-)
diff --git a/llvm/include/llvm/CodeGen/ISDOpcodes.h b/llvm/include/llvm/CodeGen/ISDOpcodes.h
index 349d1286c8dc4f..29fa3bd842c145 100644
--- a/llvm/include/llvm/CodeGen/ISDOpcodes.h
+++ b/llvm/include/llvm/CodeGen/ISDOpcodes.h
@@ -921,6 +921,8 @@ enum NodeType {
/// has native conversions.
BF16_TO_FP,
FP_TO_BF16,
+ STRICT_BF16_TO_FP,
+ STRICT_FP_TO_BF16,
/// Perform various unary floating-point operations inspired by libm. For
/// FPOWI, the result is undefined if the integer operand doesn't fit into
diff --git a/llvm/include/llvm/CodeGen/SelectionDAGNodes.h b/llvm/include/llvm/CodeGen/SelectionDAGNodes.h
index 3130f6c4dce598..d1015630b05d12 100644
--- a/llvm/include/llvm/CodeGen/SelectionDAGNodes.h
+++ b/llvm/include/llvm/CodeGen/SelectionDAGNodes.h
@@ -698,6 +698,8 @@ END_TWO_BYTE_PACK()
return false;
case ISD::STRICT_FP16_TO_FP:
case ISD::STRICT_FP_TO_FP16:
+ case ISD::STRICT_BF16_TO_FP:
+ case ISD::STRICT_FP_TO_BF16:
#define DAG_INSTRUCTION(NAME, NARG, ROUND_MODE, INTRINSIC, DAGN) \
case ISD::STRICT_##DAGN:
#include "llvm/IR/ConstrainedOps.def"
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp
index 892bfbd62f0d02..df912008dcad79 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeDAG.cpp
@@ -1034,6 +1034,7 @@ void SelectionDAGLegalize::LegalizeOp(SDNode *Node) {
Node->getOperand(0).getValueType());
break;
case ISD::STRICT_FP_TO_FP16:
+ case ISD::STRICT_FP_TO_BF16:
case ISD::STRICT_SINT_TO_FP:
case ISD::STRICT_UINT_TO_FP:
case ISD::STRICT_LRINT:
@@ -3263,12 +3264,15 @@ bool SelectionDAGLegalize::ExpandNode(SDNode *Node) {
Results.push_back(Tmp1);
break;
}
+ case ISD::STRICT_BF16_TO_FP:
+ // We don't support this expansion for now.
+ break;
case ISD::BF16_TO_FP: {
// Always expand bf16 to f32 casts, they lower to ext + shift.
//
// Note that the operand of this code can be bf16 or an integer type in case
// bf16 is not supported on the target and was softened.
- SDValue Op = Node->getOperand(0);
+ SDValue Op = Node->getOperand(Node->getOpcode() == ISD::BF16_TO_FP ? 0 : 1);
if (Op.getValueType() == MVT::bf16) {
Op = DAG.getNode(ISD::ANY_EXTEND, dl, MVT::i32,
DAG.getNode(ISD::BITCAST, dl, MVT::i16, Op));
@@ -3286,10 +3290,15 @@ bool SelectionDAGLegalize::ExpandNode(SDNode *Node) {
Results.push_back(Op);
break;
}
+ case ISD::STRICT_FP_TO_BF16:
+ // We don't support this expansion for now.
+ break;
case ISD::FP_TO_BF16: {
- SDValue Op = Node->getOperand(0);
+ bool IsStrictFP = Node->getOpcode() == ISD::STRICT_FP_TO_BF16;
+ SDValue Op = Node->getOperand(IsStrictFP ? 1 : 0);
if (Op.getValueType() != MVT::f32)
- Op = DAG.getNode(ISD::FP_ROUND, dl, MVT::f32, Op,
+ Op = DAG.getNode(IsStrictFP ? ISD::STRICT_FP_ROUND : ISD::FP_ROUND, dl,
+ MVT::f32, Op,
DAG.getIntPtrConstant(0, dl, /*isTarget=*/true));
Op = DAG.getNode(
ISD::SRL, dl, MVT::i32, DAG.getNode(ISD::BITCAST, dl, MVT::i32, Op),
@@ -4788,12 +4797,17 @@ void SelectionDAGLegalize::ConvertNodeToLibcall(SDNode *Node) {
break;
}
case ISD::STRICT_FP_EXTEND:
- case ISD::STRICT_FP_TO_FP16: {
- RTLIB::Libcall LC =
- Node->getOpcode() == ISD::STRICT_FP_TO_FP16
- ? RTLIB::getFPROUND(Node->getOperand(1).getValueType(), MVT::f16)
- : RTLIB::getFPEXT(Node->getOperand(1).getValueType(),
- Node->getValueType(0));
+ case ISD::STRICT_FP_TO_FP16:
+ case ISD::STRICT_FP_TO_BF16: {
+ RTLIB::Libcall LC = RTLIB::UNKNOWN_LIBCALL;
+ if (Node->getOpcode() == ISD::STRICT_FP_TO_FP16)
+ LC = RTLIB::getFPROUND(Node->getOperand(1).getValueType(), MVT::f16);
+ else if (Node->getOpcode() == ISD::STRICT_FP_TO_BF16)
+ LC = RTLIB::getFPROUND(Node->getOperand(1).getValueType(), MVT::bf16);
+ else
+ LC = RTLIB::getFPEXT(Node->getOperand(1).getValueType(),
+ Node->getValueType(0));
+
assert(LC != RTLIB::UNKNOWN_LIBCALL && "Unable to legalize as libcall");
TargetLowering::MakeLibCallOptions CallOptions;
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeFloatTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeFloatTypes.cpp
index f0a04589fbfdc2..ea0696be8edc4b 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeFloatTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeFloatTypes.cpp
@@ -918,6 +918,7 @@ bool DAGTypeLegalizer::SoftenFloatOperand(SDNode *N, unsigned OpNo) {
case ISD::STRICT_FP_TO_FP16:
case ISD::FP_TO_FP16: // Same as FP_ROUND for softening purposes
case ISD::FP_TO_BF16:
+ case ISD::STRICT_FP_TO_BF16:
case ISD::STRICT_FP_ROUND:
case ISD::FP_ROUND: Res = SoftenFloatOp_FP_ROUND(N); break;
case ISD::STRICT_FP_TO_SINT:
@@ -2193,13 +2194,11 @@ static ISD::NodeType GetPromotionOpcodeStrict(EVT OpVT, EVT RetVT) {
if (RetVT == MVT::f16)
return ISD::STRICT_FP_TO_FP16;
- if (OpVT == MVT::bf16) {
- // TODO: return ISD::STRICT_BF16_TO_FP;
- }
+ if (OpVT == MVT::bf16)
+ return ISD::STRICT_BF16_TO_FP;
- if (RetVT == MVT::bf16) {
- // TODO: return ISD::STRICT_FP_TO_BF16;
- }
+ if (RetVT == MVT::bf16)
+ return ISD::STRICT_FP_TO_BF16;
report_fatal_error("Attempt at an invalid promotion-related conversion");
}
@@ -2999,10 +2998,16 @@ SDValue DAGTypeLegalizer::SoftPromoteHalfRes_FP_ROUND(SDNode *N) {
EVT SVT = N->getOperand(0).getValueType();
if (N->isStrictFPOpcode()) {
- assert(RVT == MVT::f16);
- SDValue Res =
- DAG.getNode(ISD::STRICT_FP_TO_FP16, SDLoc(N), {MVT::i16, MVT::Other},
- {N->getOperand(0), N->getOperand(1)});
+ // FIXME: assume we only have two f16 variants for now.
+ unsigned Opcode;
+ if (RVT == MVT::f16)
+ Opcode = ISD::STRICT_FP_TO_FP16;
+ else if (RVT == MVT::bf16)
+ Opcode = ISD::STRICT_FP_TO_BF16;
+ else
+ llvm_unreachable("unknown half type");
+ SDValue Res = DAG.getNode(Opcode, SDLoc(N), {MVT::i16, MVT::Other},
+ {N->getOperand(0), N->getOperand(1)});
ReplaceValueWith(SDValue(N, 1), Res.getValue(1));
return Res;
}
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
index 39b7e061554141..73c7a6e4450816 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
@@ -165,6 +165,7 @@ void DAGTypeLegalizer::PromoteIntegerResult(SDNode *N, unsigned ResNo) {
case ISD::FP_TO_FP16:
Res = PromoteIntRes_FP_TO_FP16_BF16(N);
break;
+ case ISD::STRICT_FP_TO_BF16:
case ISD::STRICT_FP_TO_FP16:
Res = PromoteIntRes_STRICT_FP_TO_FP16_BF16(N);
break;
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp
index a28d834f0522f2..c0981d8362a3b0 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp
@@ -379,7 +379,9 @@ std::string SDNode::getOperationName(const SelectionDAG *G) const {
case ISD::FP_TO_FP16: return "fp_to_fp16";
case ISD::STRICT_FP_TO_FP16: return "strict_fp_to_fp16";
case ISD::BF16_TO_FP: return "bf16_to_fp";
+ case ISD::STRICT_BF16_TO_FP: return "strict_bf16_to_fp";
case ISD::FP_TO_BF16: return "fp_to_bf16";
+ case ISD::STRICT_FP_TO_BF16: return "strict_fp_to_bf16";
case ISD::LROUND: return "lround";
case ISD::STRICT_LROUND: return "strict_lround";
case ISD::LLROUND: return "llround";
More information about the llvm-commits
mailing list