Skip to content

[AMDGPU] selecting v_sat_pk instruction, version 2 #123297

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 16 commits into
base: main
Choose a base branch
from
+1,377 −69
Open

[AMDGPU] selecting v_sat_pk instruction, version 2 #123297

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
0b90a83
selecting v_sat_pk instruction, version 2
Shoreshen Jan 17, 2025
b2e29ec
fix format
Shoreshen Jan 17, 2025
4d4871c
update comments
Shoreshen Jan 17, 2025
7066fd1
upadte comments
Shoreshen Jan 17, 2025
3487eef
fix format
Shoreshen Jan 17, 2025
0cd2733
fix comments
Shoreshen Jan 17, 2025
c7c53ab
fix format
Shoreshen Jan 17, 2025
801bf90
fix comment
Shoreshen Jan 17, 2025
d89e739
update vNi8
Shoreshen Jan 21, 2025
8c51db5
add v3i16 case
Shoreshen Jan 21, 2025
c575ca5
handle N=2,4,8 case for now
Shoreshen Jan 22, 2025
b167a9b
fix test case
Shoreshen Jan 27, 2025
ebca12d
fix comments & update main
Shoreshen Jan 29, 2025
bebeb48
fix format
Shoreshen Jan 29, 2025
ef61cfb
fix comment from pierre
Shoreshen Feb 14, 2025
d178e51
Merge branch 'main' into select_v_sat_pk_v2
Shoreshen Feb 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions llvm/lib/Target/AMDGPU/AMDGPUISelLowering.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5518,6 +5518,7 @@ const char* AMDGPUTargetLowering::getTargetNodeName(unsigned Opcode) const {
NODE_NAME_CASE(UMIN3)
NODE_NAME_CASE(FMED3)
NODE_NAME_CASE(SMED3)
NODE_NAME_CASE(SAT_PK_CAST)
NODE_NAME_CASE(UMED3)
NODE_NAME_CASE(FMAXIMUM3)
NODE_NAME_CASE(FMINIMUM3)
Expand Down
1 change: 1 addition & 0 deletions llvm/lib/Target/AMDGPU/AMDGPUISelLowering.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -461,6 +461,7 @@ enum NodeType : unsigned {
FMED3,
SMED3,
UMED3,
SAT_PK_CAST,
FMAXIMUM3,
FMINIMUM3,
FDOT2,
Expand Down
3 changes: 3 additions & 0 deletions llvm/lib/Target/AMDGPU/AMDGPUInstrInfo.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -332,6 +332,9 @@ def AMDGPUumed3 : SDNode<"AMDGPUISD::UMED3", AMDGPUDTIntTernaryOp,
[]
>;

// Special node to pack v2i8 into i16 for v_sat_pk lowering.
def AMDGPUsat_pk_cast : SDNode<"AMDGPUISD::SAT_PK_CAST", SDTUnaryOp, []>;
Copy link
Contributor

@arsenm arsenm Jan 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Need to document what this is

Copy link
Contributor

@arsenm arsenm Jan 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Explain what the node is, not just to avoid v2i8. It's to pack a v2i18 into i16


def AMDGPUfmed3_impl : SDNode<"AMDGPUISD::FMED3", SDTFPTernaryOp, []>;

def AMDGPUfdot2_impl : SDNode<"AMDGPUISD::FDOT2",
Expand Down
64 changes: 64 additions & 0 deletions llvm/lib/Target/AMDGPU/SIISelLowering.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -824,6 +824,25 @@ SITargetLowering::SITargetLowering(const TargetMachine &TM,
{MVT::v4f32, MVT::v8f32, MVT::v16f32, MVT::v32f32},
Custom);
}

// True 16 instruction is current not supported
// FIXME: Add support for true 16 when supported
Copy link
Contributor

@Pierre-vh Pierre-vh Feb 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
// FIXME: Add support for true 16 when supported
// FIXME: Add support

if (!Subtarget->hasTrue16BitInsts() || !Subtarget->useRealTrue16Insts()) {
Copy link
Contributor

@arsenm arsenm Jan 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Needs predicate for has the instruction

// MVT::vNi16 for src type check in foldToSaturated
// MVT::vNi8 for dst type check in CustomLowerNode
// FIXME: Handle N = 2, 4, 8 first, should change verification logic from
// LLVM side, like break bigger vector into legal small vectors
setOperationAction(ISD::TRUNCATE_SSAT_U,
{
MVT::v2i16,
Copy link
Contributor

@arsenm arsenm Jan 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You shouldn't have to override every single type that could decompose. Ideally the combiner should be able to figure it out based on the legalizer rules

Copy link
Contributor Author

@Shoreshen Shoreshen Jan 21, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @arsenm , for truncate_ssat_u to be folded TLI.isOperationLegalOrCustom function has pass.

We didn't hook this function, so it goes default and will check getOperationAction(Op, SrcVT) == Custom, which will look up OpActions[(unsigned)VT.getSimpleVT().SimpleTy][Op], and this is set here.

If we do not set every vNi16 (source type), the related truncat_ssat_u will not be created.

Copy link
Contributor

@arsenm arsenm Jan 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For step 1 I wouldn't do this. The fix for this kind of issue is in the combiner forming them, not the legalizer rules for a specific operation

Copy link
Contributor Author

@Shoreshen Shoreshen Jan 21, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @arsenm , I agree, but changing the combiner will relate to changing llvm side code. currently I'm not planning on changing llvm side's code in this PR.

So I think we may either do this, or hook the TLI.isOperationLegalOrCustom function??

But I think hooking TLI.isOperationLegalOrCustom will make it strange, since the input variable is op code and SrcVT....

Copy link
Contributor

@arsenm arsenm Jan 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Handle the basic case in the first step, and don't worry about handling every vector perfectly right away.

Copy link
Contributor

@arsenm arsenm Feb 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You should still remove everything other than the basic case. Leave the vector splitting for a different patch

MVT::v4i16,
MVT::v8i16,
MVT::v2i8,
MVT::v4i8,
MVT::v8i8,
},
Custom);
}
}

setOperationAction({ISD::FNEG, ISD::FABS}, MVT::v4f16, Custom);
Expand Down Expand Up @@ -1983,6 +2002,12 @@ bool SITargetLowering::isTypeDesirableForOp(unsigned Op, EVT VT) const {
if (VT == MVT::i1 && Op == ISD::SETCC)
return false;

// Special case for vNi8 handling where N is even
Copy link
Contributor

@arsenm arsenm Jan 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I still don't think you should need anything in isTypeDesirableForOp

Copy link
Contributor Author

@Shoreshen Shoreshen Jan 21, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @arsenm , this function checks the destination type, while TLI.isOperationLegalOrCustom checks the source type.

The Dst type are vNi8, if we didn't return true here, it goes to the default function to check isTypeLegal(DstVT)

The backend haven't add register class for vNi8. We maybe can add the relevant register class, but makeing vNi8 legal for register class may cause unpredictable result.

Personally I think we could add the register class for relevant type when it is formally legal in the backend. So I decide to handle it here for special case.

Copy link
Contributor

@arsenm arsenm Jan 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @arsenm , this function checks the destination type, while TLI.isOperationLegalOrCustom checks the source type.

This just sounds buggy. The interpretation of which type is the one that matters for the opcode needs to be globally consistent

The backend haven't add register class for vNi8. We maybe can add the relevant register class,

Please no, this is a huge amount of work.

Copy link
Contributor Author

@Shoreshen Shoreshen Jan 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @arsenm , yeah it is kind of weird logic. And what made it more strange is that to get into the ReplaceNodeResults function, it will check TLI.getOperationAction(Opc, DstVT) == Custom........

But if we want to change this, I think we also need to modify AArch64 backend....

Copy link
Contributor

@arsenm arsenm Jan 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, bugs cause other bugs and all the use points need to be fixed

Copy link
Contributor

@arsenm arsenm Feb 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I still do not think this is necessary in this patch and should be dropped

if (Op == ISD::TRUNCATE_SSAT_U && VT.isVector() &&
VT.getVectorElementType() == MVT::i8 &&
((VT.getVectorNumElements() & 1) == 0))
return true;

return TargetLowering::isTypeDesirableForOp(Op, VT);
}

Expand Down Expand Up @@ -6615,6 +6640,45 @@ void SITargetLowering::ReplaceNodeResults(SDNode *N,
Results.push_back(lowerFSQRTF16(SDValue(N, 0), DAG));
break;
}
case ISD::TRUNCATE_SSAT_U: {
SDLoc SL(N);
SDValue Src = N->getOperand(0);
EVT SrcVT = Src.getValueType();
EVT DstVT = N->getValueType(0);

assert(SrcVT.isVector() && DstVT.isVector());
Copy link
Contributor

@arsenm arsenm Jan 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also should assert the element type is i8

assert(DstVT.getVectorElementType() == MVT::i8);
assert(SrcVT.getVectorElementType() == MVT::i16);

unsigned EleNo = SrcVT.getVectorNumElements();
assert(EleNo == DstVT.getVectorNumElements());

if (EleNo == 2) {
SDValue Op = DAG.getNode(AMDGPUISD::SAT_PK_CAST, SL, MVT::i16, Src);
Op = DAG.getNode(ISD::BITCAST, SL, N->getValueType(0), Op);
Results.push_back(Op);
break;
}

// Vector case, number of element must be even
Copy link
Contributor

@arsenm arsenm Feb 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Leave this for second patch, I would hope the legalization process ends up with this in the easy form

assert((EleNo & 1) == 0);
SmallVector<SDValue> DstPairs;
EVT SrcEleVT = SrcVT.getVectorElementType();
EVT DstEleVT = DstVT.getVectorElementType();
EVT SrcPairVT = EVT::getVectorVT(*DAG.getContext(), SrcEleVT, 2);
EVT DstPairVT = EVT::getVectorVT(*DAG.getContext(), DstEleVT, 2);
for (unsigned i = 0; i != EleNo; i += 2) {
SDValue SrcPair = DAG.getNode(ISD::EXTRACT_SUBVECTOR, SL, SrcPairVT, Src,
DAG.getConstant(i, SL, MVT::i32));
SDValue SatPk =
DAG.getNode(AMDGPUISD::SAT_PK_CAST, SL, MVT::i16, SrcPair);
SDValue DstPair = DAG.getNode(ISD::BITCAST, SL, DstPairVT, SatPk);
DstPairs.push_back(DstPair);
}
SDValue Op = DAG.getNode(ISD::CONCAT_VECTORS, SL, DstVT, DstPairs);
Results.push_back(Op);
Copy link
Contributor

@arsenm arsenm Jan 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You need to insert a cast to the original type, does this not assert as-is?

break;
}
default:
AMDGPUTargetLowering::ReplaceNodeResults(N, Results, DAG);
break;
Expand Down
15 changes: 15 additions & 0 deletions llvm/lib/Target/AMDGPU/SIInstructions.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3382,6 +3382,21 @@ def : GCNPat <
(v2i16 (V_LSHL_OR_B32_e64 $src1, (i32 16), (i32 (V_AND_B32_e64 (i32 (V_MOV_B32_e32 (i32 0xffff))), $src0))))
>;

multiclass V_SAT_PK_Pat<Instruction inst> {
def : GCNPat<
Copy link
Contributor

@arsenm arsenm Jan 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This pattern isn't doing much, you should be able to pass the node to the SDNodeOperator argument to the instruction definition

Copy link
Contributor Author

@Shoreshen Shoreshen Jan 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @arsenm , could you be more specific? Should I use other type of pattern?

Copy link
Contributor

@arsenm arsenm Jan 17, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The node is basically the same as the instruction definition, so you should be able to use the built-in pattern attached to the instruction def.
Something like

in VOP1Instructions.td:
defm V_SAT_PK_U8_I16 : VOP1Inst_t16<"v_sat_pk_u8_i16", VOP_I16_I32, AMDGPUsat_pk_cast>;

Copy link
Contributor Author

@Shoreshen Shoreshen Jan 17, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @arsenm , by adding the node I got the following:

def V_SAT_PK_U8_I16_e64: list<dag> Pattern = [(set i16:$vdst, (AMDGPUsat_pk_cast (i32 (VOP3Mods0 i32:$src0))))];
def V_SAT_PK_U8_I16_fake16_e64: list<dag> Pattern = [(set i16:$vdst, (AMDGPUsat_pk_cast (i32 (VOP3Mods0 i32:$src0))))];
def V_SAT_PK_U8_I16_t16_e64: list<dag> Pattern = [(set i16:$vdst, (AMDGPUsat_pk_cast (i32 (VOP3OpSelMods i32:$src0, i32:$src0_modifiers))))];

I think there are 2 problems:

  1. The source is i32, instead of v2i16
  2. It requires the operand of AMDGPUsat_pk_cast be complex pattern of VOP3Mods0 or VOP3OpSelMods

If the instruction cannot cover any type of (i16 (AMDGPUsat_pk_cast v2i8)), we gain risk of failing in selection.

I also tried to create a new VOP_I16_V2I16 type, but it makes V_SAT_PK_U8_I16_e64 and V_SAT_PK_U8_I16_fake16_e64 4 operands instructions (with modifier, clamp and opsel)

I think in order to make the passing node work, I need to modify related complex pattern functions and replace (v2i8 (truncssat_u v2i16)) with some patterns that can fit the complex pattern functions

(i16 (AMDGPUsat_pk_cast v2i16:$src)),
(inst VRegSrc_32:$src)
>;
}

let OtherPredicates = [NotHasTrue16BitInsts] in {
defm : V_SAT_PK_Pat<V_SAT_PK_U8_I16_e64>;
} // End OtherPredicates = [NotHasTrue16BitInsts]

let True16Predicate = UseFakeTrue16Insts in {
defm : V_SAT_PK_Pat<V_SAT_PK_U8_I16_fake16_e64>;
} // End True16Predicate = UseFakeTrue16Insts

// With multiple uses of the shift, this will duplicate the shift and
// increase register pressure.
def : GCNPat <
Expand Down
Loading
Loading