Nvptx port LowerBITCAST to SelectionDAG #120903

GrumpyPigSkin · 2024-12-22T15:20:12Z

Addressing #116695: Ports Bitcast to SelectionDAG

…zeTypes.cpp.

llvmbot · 2024-12-22T15:20:42Z

@llvm/pr-subscribers-llvm-selectiondag

@llvm/pr-subscribers-backend-nvptx

Author: None (GrumpyPigSkin)

Changes

Addressing #116695: Ports Bitcast to SelectionDAG

Full diff: https://github.com/llvm/llvm-project/pull/120903.diff

5 Files Affected:

(modified) llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp (+6-2)
(modified) llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp (+32)
(modified) llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h (+1)
(modified) llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp (+1-25)
(modified) llvm/lib/Target/NVPTX/NVPTXISelLowering.h (-2)

diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
index be7521f3416850..bcb59e3c2aef3e 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
@@ -2174,8 +2174,12 @@ SDValue DAGTypeLegalizer::PromoteIntOp_ATOMIC_STORE(AtomicSDNode *N) {
 }
 
 SDValue DAGTypeLegalizer::PromoteIntOp_BITCAST(SDNode *N) {
-  // This should only occur in unusual situations like bitcasting to an
-  // x86_fp80, so just turn it into a store+load
+  // Use the custom lowering.
+  if (const auto Res = LowerBitcast(N)) {
+    return Res;
+  }
+
+  // If it fails fall back to the default method
   return CreateStackStoreLoad(N->getOperand(0), N->getValueType(0));
 }
 
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp
index b6abad830c371e..4aecf667b2cee1 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp
@@ -910,6 +910,38 @@ SDValue DAGTypeLegalizer::CreateStackStoreLoad(SDValue Op,
   return DAG.getLoad(DestVT, dl, Store, StackPtr, MachinePointerInfo(), Align);
 }
 
+static SDValue MaybeBitcast(SelectionDAG &DAG, SDLoc DL, EVT VT,
+                            SDValue Value) {
+  if (Value->getValueType(0) == VT)
+    return Value;
+  return DAG.getNode(ISD::BITCAST, DL, VT, Value);
+}
+
+SDValue DAGTypeLegalizer::LowerBitcast(SDNode *Node) const {
+  assert(Node->getOpcode() == ISD::BITCAST && "Unexpected opcode!");
+  // Handle bitcasting from v2i8 without hitting the default promotion
+  // strategy which goes through stack memory.
+  EVT FromVT = Node->getOperand(0)->getValueType(0);
+  if (FromVT != MVT::v2i8) {
+    return SDValue();
+  }
+
+  // Pack vector elements into i16 and bitcast to final type
+  SDLoc DL(Node);
+  SDValue Vec0 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, MVT::i8,
+                             Node->getOperand(0), DAG.getIntPtrConstant(0, DL));
+  SDValue Vec1 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, MVT::i8,
+                             Node->getOperand(0), DAG.getIntPtrConstant(1, DL));
+  SDValue Extend0 = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i16, Vec0);
+  SDValue Extend1 = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i16, Vec1);
+  SDValue Const8 = DAG.getConstant(8, DL, MVT::i16);
+  SDValue AsInt = DAG.getNode(
+      ISD::OR, DL, MVT::i16,
+      {Extend0, DAG.getNode(ISD::SHL, DL, MVT::i16, {Extend1, Const8})});
+  EVT ToVT = Node->getValueType(0);
+  return MaybeBitcast(DAG, DL, ToVT, AsInt);
+}
+
 /// Replace the node's results with custom code provided by the target and
 /// return "true", or do nothing and return "false".
 /// The last parameter is FALSE if we are dealing with a node with legal
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
index 571a710cc92a34..30951112069ed5 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
@@ -216,6 +216,7 @@ class LLVM_LIBRARY_VISIBILITY DAGTypeLegalizer {
   SDValue BitConvertToInteger(SDValue Op);
   SDValue BitConvertVectorToIntegerVector(SDValue Op);
   SDValue CreateStackStoreLoad(SDValue Op, EVT DestVT);
+  SDValue LowerBitcast(SDNode *N) const;
   bool CustomLowerNode(SDNode *N, EVT VT, bool LegalizeResult);
   bool CustomWidenLowerNode(SDNode *N, EVT VT);
 
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
index 5c1f717694a4c7..2eaeb624004730 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
@@ -2086,30 +2086,6 @@ NVPTXTargetLowering::LowerCONCAT_VECTORS(SDValue Op, SelectionDAG &DAG) const {
   return DAG.getBuildVector(Node->getValueType(0), dl, Ops);
 }
 
-SDValue NVPTXTargetLowering::LowerBITCAST(SDValue Op, SelectionDAG &DAG) const {
-  // Handle bitcasting from v2i8 without hitting the default promotion
-  // strategy which goes through stack memory.
-  EVT FromVT = Op->getOperand(0)->getValueType(0);
-  if (FromVT != MVT::v2i8) {
-    return Op;
-  }
-
-  // Pack vector elements into i16 and bitcast to final type
-  SDLoc DL(Op);
-  SDValue Vec0 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, MVT::i8,
-                             Op->getOperand(0), DAG.getIntPtrConstant(0, DL));
-  SDValue Vec1 = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, MVT::i8,
-                             Op->getOperand(0), DAG.getIntPtrConstant(1, DL));
-  SDValue Extend0 = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i16, Vec0);
-  SDValue Extend1 = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i16, Vec1);
-  SDValue Const8 = DAG.getConstant(8, DL, MVT::i16);
-  SDValue AsInt = DAG.getNode(
-      ISD::OR, DL, MVT::i16,
-      {Extend0, DAG.getNode(ISD::SHL, DL, MVT::i16, {Extend1, Const8})});
-  EVT ToVT = Op->getValueType(0);
-  return MaybeBitcast(DAG, DL, ToVT, AsInt);
-}
-
 // We can init constant f16x2/v2i16/v4i8 with a single .b32 move.  Normally it
 // would get lowered as two constant loads and vector-packing move.
 // Instead we want just a constant move:
@@ -2619,7 +2595,7 @@ NVPTXTargetLowering::LowerOperation(SDValue Op, SelectionDAG &DAG) const {
   case ISD::BUILD_VECTOR:
     return LowerBUILD_VECTOR(Op, DAG);
   case ISD::BITCAST:
-    return LowerBITCAST(Op, DAG);
+    return SDValue();
   case ISD::EXTRACT_SUBVECTOR:
     return Op;
   case ISD::EXTRACT_VECTOR_ELT:
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.h b/llvm/lib/Target/NVPTX/NVPTXISelLowering.h
index 4a98fe21b81dc6..446ff1536d36cf 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.h
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.h
@@ -265,8 +265,6 @@ class NVPTXTargetLowering : public TargetLowering {
   const NVPTXSubtarget &STI; // cache the subtarget here
   SDValue getParamSymbol(SelectionDAG &DAG, int idx, EVT) const;
 
-  SDValue LowerBITCAST(SDValue Op, SelectionDAG &DAG) const;
-
   SDValue LowerBUILD_VECTOR(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerCONCAT_VECTORS(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerEXTRACT_VECTOR_ELT(SDValue Op, SelectionDAG &DAG) const;

GrumpyPigSkin · 2024-12-22T15:34:13Z

@justinfargnoli Could you please review this, I have a couple of questions/notes also:

I have put MaybeBitcast in LegalizeTypes.cpp but there is also an exact copy in NVPTXISelLowering.cpp is there some common place this could go to remove the duplication?
I wasn't 100% sure which class the LowerBitcast method should be defined in so if there is a more suitable place lmk.
I added the operation inside of PromoteIntOp_BITCAST but again, I wasn't 100% sure if this was the best place to apply the transform. If CustomLowerNode fails to lower, it does the generic transform inside of PromoteIntOp_BITCAST, so I thought that could be a good place to do it.

Thanks :)

RKSimon

A few minors

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

RKSimon · 2024-12-22T16:15:51Z

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

+  EVT FromVT = Node->getOperand(0)->getValueType(0);
+  if (FromVT != MVT::v2i8) {
+    return SDValue();
+  }


(style) remove braces

Are we gaining anything from having PackBitcastInRegister/UnpackBitcastInRegister vs merging them inside LowerBitcastInRegister?

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

RKSimon · 2024-12-22T16:17:34Z

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

+                             Node->getOperand(0), DAG.getIntPtrConstant(1, DL));
+  SDValue Extend0 = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i16, Vec0);
+  SDValue Extend1 = DAG.getNode(ISD::ZERO_EXTEND, DL, MVT::i16, Vec1);
+  SDValue Const8 = DAG.getConstant(8, DL, MVT::i16);


Use DAG.getShiftAmountConstant() - some targets don't have matching shift value / amount operand types.

Could you please check that what I have put is correct? I followed other examples of using getShiftAmountConstant

RKSimon · 2024-12-22T16:19:04Z

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

    return LowerBUILD_VECTOR(Op, DAG);
  case ISD::BITCAST:
-    return LowerBITCAST(Op, DAG);
+    return SDValue();


Can you remove this case entirely? You might need to remove any custom setOperationAction entries for BITCAST?

I removed the case and the setOperationAction but when I did that it started failing the second test case and applying the stack-based bitcast. I am unsure why it did that so I will look into it further. Unless you have any ideas?

DAGTypeLegalizer::ExpandRes_BITCAST already does what you're after (expand bitcast of vector to scalar integer) in a more general form - maybe start there?

Hi, I had a look at DAGTypeLegalizer::ExpandRes_BITCAST it is close to what I want but it emits the store/load.
I managed to get the code working for the NVPTX test case by removing the setOperationAction. Please will look at the newer code and see if the idea is on the right track?

github-actions · 2024-12-22T16:58:25Z

✅ With the latest revision this PR passed the C/C++ code formatter.

arsenm · 2024-12-23T07:06:37Z

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

+  EVT ShiftAmtTy =
+      TLI.getShiftAmountTy(Extend1.getValueType(), DAG.getDataLayout());
+  SDValue ShiftConst = DAG.getShiftAmountConstant(8, ShiftAmtTy, DL);


This isn't correct usage. getShiftAmountConstant internally uses getShiftAmountTy. If you use getShiftAmountTy, you use the type directly to getConstant. Or you could use getShiftAmountConstant with Extend1.getValueType()

Thanks :) I have updated the code

arsenm · 2024-12-23T07:07:13Z

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

+  // Handle bitcasting from v2i8 without hitting the default promotion
+  // strategy which goes through stack memory.
+  EVT FromVT = Node->getOperand(0)->getValueType(0);
+  if (FromVT != MVT::v2i8)


This is too specific of a type check for the generic legalizer. This code should be generalized over the types supported

What other types do you think we should support? Or would this be the preferred method from int->vec and vice versa? ie.

<2 x i8> to i16 <4 x i8> to i32 <2 x i16> to i32 ; etc

I could see that changing a lot of code though.

RKSimon

Just finishing my last day before Christmas break but I'll try to finish my review when I can.

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

GrumpyPigSkin · 2025-01-16T19:20:51Z

@RKSimon with the code as it is now, it passes the NVPTX test that was applying this lowering before. However, because the function is generalised it results in 25 test failures due to the difference in code when performing bitcasts. I presume we don't want code to change across the board? Is there a limit to the size of the types we want to shift in register, and are there any platforms we don't want to apply it to?

RKSimon · 2025-01-20T11:08:53Z

You'd be better off including all the test changes in the PR instead of having the CI fail - then reviewers can better advise on what you need to do

GrumpyPigSkin · 2025-01-20T23:23:17Z

@RKSimon I have update all the failing tests using the update_llc_test_checks.py script :)

arsenm · 2025-01-22T07:02:48Z

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

+
+    return DAG.getBitcast(ToVT, Packed);
+
+  } else if (FromVT.isScalarInteger() && ToVT.isVector()) {


No else after return

arsenm · 2025-01-22T07:03:25Z

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

+      unsigned ElementIndex = IsBigEndian ? (NumElems - 1 - I) : I;
+      SDValue Elem =
+          DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, ElemVT, N->getOperand(0),
+                      DAG.getIntPtrConstant(ElementIndex, DL));


This should be using getVectorIdxTy

arsenm · 2025-01-22T07:04:29Z

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

+      unsigned ShiftAmountVal = ElemBits * ElementIndex;
+
+      SDValue ShiftAmount =
+          DAG.getShiftAmountConstant(ShiftAmountVal, FromVT, DL);


Pull the getShiftAmountTy out of the loop? Can then use a shorter getConstant inside

arsenm · 2025-01-22T07:05:22Z

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

+      SDValue Shifted =
+          DAG.getNode(ISD::SRL, DL, FromVT, N->getOperand(0), ShiftAmount);
+      SDValue Element = DAG.getNode(ISD::TRUNCATE, DL, ElemVT, Shifted);
+      Elements.push_back(Element);


Size the vector at the start and assign the element at index

RKSimon · 2025-01-24T18:10:05Z

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

+    assert(PackedBits >= ElemBits * NumElems &&
+           "Scalar type does not have enough bits to pack vector values.");
+
+    EVT PackVT = EVT::getIntegerVT(*DAG.getContext(), ElemBits * NumElems);


How can PackVT != ToVT if its a scalar integer?

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp

RKSimon · 2025-02-04T10:23:42Z

@arsenm Any suggestions on how to continue with this after #125637 ?

arsenm · 2025-02-05T12:55:06Z

@arsenm Any suggestions on how to continue with this after #125637 ?

There are plenty of cases that still hit the stack path, that shouldn't. I don't think this conflicts that much

GrumpyPigSkin added 2 commits December 22, 2024 15:08

Ported LowerBITCAST from NVPTXISelLowering.cpp to SelectionDAG/Legali…

e1b6fce

…zeTypes.cpp.

Removed redundant assert check

3320d55

llvmbot added backend:NVPTX llvm:SelectionDAG SelectionDAGISel as well labels Dec 22, 2024

RKSimon reviewed Dec 22, 2024

View reviewed changes

Addressed Most Code Review Comments

d7cb133

Applied code formatting

b674487

arsenm reviewed Dec 23, 2024

View reviewed changes

GrumpyPigSkin added 2 commits December 23, 2024 17:23

Generalised bit packing and unpacking

f8dadde

Formatting

a6ca08a

RKSimon reviewed Dec 23, 2024

View reviewed changes

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp Outdated Show resolved Hide resolved

Check for only scalar integer types

75a533c

GrumpyPigSkin requested a review from RKSimon January 13, 2025 19:12

GrumpyPigSkin added 3 commits January 16, 2025 18:50

Merged function

686f698

Formatting

d8a82eb

Updated comments

e5f7304

Updated failing tests

8d01cbe

arsenm reviewed Jan 22, 2025

View reviewed changes

GrumpyPigSkin added 2 commits January 22, 2025 20:13

Code review fixes

87530e4

Applied formatting

7a5cfe4

RKSimon reviewed Jan 24, 2025

View reviewed changes

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.cpp Show resolved Hide resolved

GrumpyPigSkin added 2 commits January 25, 2025 15:47

Addressed code review comments

66bfc99

Formatting

f278d56


		return DAG.getBitcast(ToVT, Packed);

		} else if (FromVT.isScalarInteger() && ToVT.isVector()) {

Nvptx port LowerBITCAST to SelectionDAG #120903

Are you sure you want to change the base?

Nvptx port LowerBITCAST to SelectionDAG #120903

Uh oh!

Conversation

GrumpyPigSkin commented Dec 22, 2024

Uh oh!

llvmbot commented Dec 22, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

GrumpyPigSkin commented Dec 22, 2024

Uh oh!

RKSimon left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Dec 22, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

RKSimon left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

GrumpyPigSkin commented Jan 16, 2025

Uh oh!

RKSimon commented Jan 20, 2025

Uh oh!

GrumpyPigSkin commented Jan 20, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

RKSimon commented Feb 4, 2025

Uh oh!

arsenm commented Feb 5, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

llvmbot commented Dec 22, 2024 •

edited

Loading

github-actions bot commented Dec 22, 2024 •

edited

Loading