Skip to content

[RISCV][GlobalISel] Legalize Scalable Vector Loads and Stores #84965

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.

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

Merged
merged 11 commits into from
Jul 31, 2024

Conversation

jiahanxie353
Copy link
Contributor

@jiahanxie353 jiahanxie353 commented Mar 12, 2024

This patch works on legalizing load and store instruction for scalable vectors

Copy link

github-actions bot commented Mar 25, 2024

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

Copy link

✅ With the latest revision this PR passed the Python code formatter.

@michaelmaitland michaelmaitland self-requested a review March 25, 2024 16:24
@michaelmaitland
Copy link
Contributor

Can you please update PR description?

@llvmbot
Copy link
Member

llvmbot commented Mar 25, 2024

@llvm/pr-subscribers-llvm-globalisel

@llvm/pr-subscribers-backend-risc-v

Author: Jiahan Xie (jiahanxie353)

Changes

This patch works on legalizing load instruction for scalable vectors


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

5 Files Affected:

  • (modified) llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h (+2-1)
  • (modified) llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp (+6-2)
  • (modified) llvm/lib/CodeGen/MIRParser/MIParser.cpp (+1-1)
  • (modified) llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp (+6-6)
  • (added) llvm/test/CodeGen/RISCV/GlobalISel/legalizer/rvv/legalize-load.mir (+29)
diff --git a/llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h b/llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h
index 6afaea3f3fc5c6..5d60f4f1829397 100644
--- a/llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h
+++ b/llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h
@@ -220,7 +220,8 @@ struct TypePairAndMemDesc {
            Align >= Other.Align &&
            // FIXME: This perhaps should be stricter, but the current legality
            // rules are written only considering the size.
-           MemTy.getSizeInBits() == Other.MemTy.getSizeInBits();
+           MemTy.getSizeInBits().getKnownMinValue() ==
+               Other.MemTy.getSizeInBits().getKnownMinValue();
   }
 };
 
diff --git a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
index 1b25da8833e4fb..fb18a2fd4d3e8c 100644
--- a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
@@ -3330,8 +3330,12 @@ LegalizerHelper::LegalizeResult LegalizerHelper::lowerLoad(GAnyLoad &LoadMI) {
   LLT MemTy = MMO.getMemoryType();
   MachineFunction &MF = MIRBuilder.getMF();
 
-  unsigned MemSizeInBits = MemTy.getSizeInBits();
-  unsigned MemStoreSizeInBits = 8 * MemTy.getSizeInBytes();
+  unsigned MemSizeInBits = MemTy.isScalable()
+                               ? MemTy.getSizeInBits().getKnownMinValue()
+                               : MemTy.getSizeInBits();
+  unsigned MemStoreSizeInBits =
+      MemTy.isScalable() ? 8 * MemTy.getSizeInBytes().getKnownMinValue()
+                         : 8 * MemTy.getSizeInBytes();
 
   if (MemSizeInBits != MemStoreSizeInBits) {
     if (MemTy.isVector())
diff --git a/llvm/lib/CodeGen/MIRParser/MIParser.cpp b/llvm/lib/CodeGen/MIRParser/MIParser.cpp
index 691c60d22724f3..43f6d3219bc6da 100644
--- a/llvm/lib/CodeGen/MIRParser/MIParser.cpp
+++ b/llvm/lib/CodeGen/MIRParser/MIParser.cpp
@@ -3415,7 +3415,7 @@ bool MIParser::parseMachineMemoryOperand(MachineMemOperand *&Dest) {
     if (expectAndConsume(MIToken::rparen))
       return true;
 
-    Size = MemoryType.getSizeInBytes();
+    Size = MemoryType.getSizeInBytes().getKnownMinValue();
   }
 
   MachinePointerInfo Ptr = MachinePointerInfo();
diff --git a/llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp b/llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp
index 64ae4e94a8c929..bcded4b227e7ef 100644
--- a/llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp
+++ b/llvm/lib/Target/RISCV/GISel/RISCVLegalizerInfo.cpp
@@ -210,12 +210,12 @@ RISCVLegalizerInfo::RISCVLegalizerInfo(const RISCVSubtarget &ST)
       .clampScalar(0, s32, (XLen == 64 || ST.hasStdExtD()) ? s64 : s32)
       .clampScalar(1, sXLen, sXLen);
 
-  auto &LoadStoreActions =
-      getActionDefinitionsBuilder({G_LOAD, G_STORE})
-          .legalForTypesWithMemDesc({{s32, p0, s8, 8},
-                                     {s32, p0, s16, 16},
-                                     {s32, p0, s32, 32},
-                                     {p0, p0, sXLen, XLen}});
+  auto &LoadStoreActions = getActionDefinitionsBuilder({G_LOAD, G_STORE})
+                               .legalForTypesWithMemDesc({{s32, p0, s8, 8},
+                                                          {s32, p0, s16, 16},
+                                                          {s32, p0, s32, 32},
+                                                          {p0, p0, sXLen, XLen},
+                                                          {nxv1s8, p0, s8, 8}});
   auto &ExtLoadActions =
       getActionDefinitionsBuilder({G_SEXTLOAD, G_ZEXTLOAD})
           .legalForTypesWithMemDesc({{s32, p0, s8, 8}, {s32, p0, s16, 16}});
diff --git a/llvm/test/CodeGen/RISCV/GlobalISel/legalizer/rvv/legalize-load.mir b/llvm/test/CodeGen/RISCV/GlobalISel/legalizer/rvv/legalize-load.mir
new file mode 100644
index 00000000000000..1b62f555207165
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/GlobalISel/legalizer/rvv/legalize-load.mir
@@ -0,0 +1,29 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=riscv32 -mattr=+v -run-pass=legalizer %s -o - | FileCheck %s
+# RUN: llc -mtriple=riscv64 -mattr=+v -run-pass=legalizer %s -o - | FileCheck %s
+--- |
+
+  define <vscale x 1 x i8> @vload_nx1i8(ptr %pa) {
+    %va = load <vscale x 1 x i8>, ptr %pa
+    ret <vscale x 1 x i8> %va
+  }
+
+...
+---
+name:            vload_nx1i8
+body:             |
+  bb.1 (%ir-block.0):
+    liveins: $x10
+
+    ; CHECK-LABEL: name: vload_nx1i8
+    ; CHECK: liveins: $x10
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(p0) = COPY $x10
+    ; CHECK-NEXT: [[LOAD:%[0-9]+]]:_(<vscale x 1 x s8>) = G_LOAD [[COPY]](p0) :: (load (<vscale x 1 x s8>) from %ir.pa)
+    ; CHECK-NEXT: $v8 = COPY [[LOAD]](<vscale x 1 x s8>)
+    ; CHECK-NEXT: PseudoRET implicit $v8
+    %0:_(p0) = COPY $x10
+    %1:_(<vscale x 1 x s8>) = G_LOAD %0(p0) :: (load (<vscale x 1 x s8>) from %ir.pa)
+    $v8 = COPY %1(<vscale x 1 x s8>)
+    PseudoRET implicit $v8
+

@jiahanxie353 jiahanxie353 changed the title [RISCV][GlobalISel] Legalize Scalable Vector Loads [RISCV][GlobalISel] Legalize Scalable Vector Loads and Stores Mar 26, 2024
@@ -3330,16 +3330,17 @@ LegalizerHelper::LegalizeResult LegalizerHelper::lowerLoad(GAnyLoad &LoadMI) {
LLT MemTy = MMO.getMemoryType();
MachineFunction &MF = MIRBuilder.getMF();

unsigned MemSizeInBits = MemTy.getSizeInBits();
unsigned MemStoreSizeInBits = 8 * MemTy.getSizeInBytes();
unsigned MinMemSizeInBits = MemTy.getSizeInBits().getKnownMinValue();
Copy link
Contributor

Choose a reason for hiding this comment

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

I think we can do this as TypeSize MemSizeInBits and TypeSize MinMemStoreSizeInBits as Craig Pointed out, which allows us to avoid calling getKnownMinValue. That way we avoid the case comparing fixed vector and scalable vector.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Understood.
I used getKnownMinValue because we have cases for doing scalar arithmetic operations below. For example:

if (!isPowerOf2_32(MinMemSizeInBits)) {
// This load needs splitting into power of 2 sized loads.
LargeSplitSize = llvm::bit_floor(MinMemSizeInBits);
SmallSplitSize = MinMemSizeInBits - LargeSplitSize;
} else {
// This is already a power of 2, but we still need to split this in half.
//
// Assume we're being asked to decompose an unaligned load.
// TODO: If this requires multiple splits, handle them all at once.
auto &Ctx = MF.getFunction().getContext();
if (TLI.allowsMemoryAccess(Ctx, MIRBuilder.getDataLayout(), MemTy, MMO))
return UnableToLegalize;
SmallSplitSize = LargeSplitSize = MinMemSizeInBits / 2;
}

bit_floor has to take a scalar.

So are you suggesting we make them TypeSize MemSizeInBits but call getKnownMinValue on demand?

@jiahanxie353 jiahanxie353 force-pushed the vload-legalize branch 2 times, most recently from 699b864 to aab0ab4 Compare March 27, 2024 18:08
@jiahanxie353
Copy link
Contributor Author

When we have alignment restrictions for load, we need to check allowsMemoryAccessForAlignment, like what SelectionDAG did. Therefore, equivalently, we can insert customIf(!allowsMemoryAccessForAlignment) around here

But to get such information, we need to have LLVMContext, DataLayout, etc. However, RISCVLegalizerInfo does not have those fields in it. Nor can I figure out a way to use LegalityQuery to obtain those required arguments.

Can I get some help on this? Thanks!

@michaelmaitland
Copy link
Contributor

When we have alignment restrictions for load, we need to check allowsMemoryAccessForAlignment, like what SelectionDAG did. Therefore, equivalently, we can insert customIf(!allowsMemoryAccessForAlignment) around here

But to get such information, we need to have LLVMContext, DataLayout, etc. However, RISCVLegalizerInfo does not have those fields in it. Nor can I figure out a way to use LegalityQuery to obtain those required arguments.

Can I get some help on this? Thanks!

Check out legalizeIntrinsic in RISCVLegalizerInfo. You can get the DataLayout from MIRBuilder.getDataLayout(). You can get LLVMContext from MF.getFunction.getContext().

These are all available in legalizeCustom but not in the LegalityQuery, since the LegalityQuery gets built before seeing any instructions or functions. That means you will need to write a customIf that sends us to legalizeCustom more circumstances than are actually legal -- we will likely have to return false in some instances from legalizeCustom.

@topperc
Copy link
Collaborator

topperc commented May 28, 2024

When we have alignment restrictions for load, we need to check allowsMemoryAccessForAlignment, like what SelectionDAG did. Therefore, equivalently, we can insert customIf(!allowsMemoryAccessForAlignment) around here

But to get such information, we need to have LLVMContext, DataLayout, etc. However, RISCVLegalizerInfo does not have those fields in it. Nor can I figure out a way to use LegalityQuery to obtain those required arguments.

Can I get some help on this? Thanks!

The DataLayout belongs to Module. The LegalizerInfo constructor runs before the Module exists or at least without knowledge of the Module. I don't think you can call allowsMemoryAccessForAlignment.

A misaligned vector can be handled by converting it to a vle8. Not sure if the generic legalization can handle that or we need to do custom legalization for RISC-V. SelectionDAG uses custom legalization, see RISCVTargetLowering::expandUnalignedRVVLoad

}

define <vscale x 16 x i8> @vload_nx16i8(ptr %pa) {
%va = load <vscale x 16 x i8>, ptr %pa, align 16
Copy link
Contributor

Choose a reason for hiding this comment

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

should these be store instructions?


DstOp NewDstReg(NewLoadTy);
MachineMemOperand *NewLoadMMO = MF->getMachineMemOperand(
MachinePointerInfo(), MachineMemOperand::MOLoad, NewLoadTy, Alignment);
Copy link
Contributor

Choose a reason for hiding this comment

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

Should we be casting the original MI to GLoad and building the MachineMemOperand using existing info. It looks like RISCVTargetLowering::ExpandUnalignedRVVLoad uses the pointer info from the original load.

Copy link
Contributor

Choose a reason for hiding this comment

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

MachinePointerInfo PI = cast<GLoad>(MI)->getMMO()->getPointerInfo();
MachineMemOperand *NewLoadMMO = MF->getMachineMemOperand(PI, MachineMemOperand::MOLoad, NewLoadTy, Alignment)

"Load instructions only have one MemOperand.");
Align Alignment = (*MI.memoperands_begin())->getAlign();
MachineMemOperand *LoadMMO = MF->getMachineMemOperand(
MachinePointerInfo(), MachineMemOperand::MOLoad, LoadTy, Alignment);
Copy link
Contributor

Choose a reason for hiding this comment

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

Should we be casting the original MI to GLoad and building the MachineMemOperand using existing info. It looks like RISCVTargetLowering::ExpandUnalignedRVVLoad uses the mem operand from the original load.


auto NewLoad = MIB.buildLoad(NewDstReg, PtrReg, *NewLoadMMO);

MIB.buildBitcast(DstReg, NewLoad.getReg(0));
Copy link
Contributor

Choose a reason for hiding this comment

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

I think you might be able to pass NewLoad here instead of NewLoad.getReg(0).

MachineMemOperand *NewLoadMMO = MF->getMachineMemOperand(
MachinePointerInfo(), MachineMemOperand::MOLoad, NewLoadTy, Alignment);

auto NewLoad = MIB.buildLoad(NewDstReg, PtrReg, *NewLoadMMO);
Copy link
Contributor

Choose a reason for hiding this comment

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

I think you are using this function for both load and store but the logic in it is only assuming it is legalizing a load.

Comment on lines +708 to +735
if (TLI->allowsMemoryAccessForAlignment(Ctx, DL, VT, *MMO))
return true;
Copy link
Contributor

Choose a reason for hiding this comment

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

This should really be considered in the default lower action. I've had a patch for years I need to get back to


LoadStoreActions.widenScalarToNextPow2(0, /* MinSize = */ 8)
.lowerIfMemSizeNotByteSizePow2()
.customIf(all(LegalityPredicate([=](const LegalityQuery &Query) {
Copy link
Contributor

Choose a reason for hiding this comment

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

nit: IIUC the customIf is to handle vector types that have non-standard align arguments. If that is an accurate description, could you add a comment clarifying this point for the reader? If someone isn't familiar with the code, it might not be clear why we are doing this special custom legalization for vector types that we addressed above.

@@ -3690,7 +3690,7 @@ LegalizerHelper::LegalizeResult LegalizerHelper::lowerLoad(GAnyLoad &LoadMI) {

if (!isPowerOf2_32(MemSizeInBits)) {
// This load needs splitting into power of 2 sized loads.
LargeSplitSize = llvm::bit_floor(MemSizeInBits);
LargeSplitSize = llvm::bit_floor(MemSizeInBits.getKnownMinValue());
Copy link
Contributor

Choose a reason for hiding this comment

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

These values are used in the call to getMachineMemOperand. There is a signature:

MachineMemOperand *MachineFunction::getMachineMemOperand(
    const MachineMemOperand *MMO, const MachinePointerInfo &PtrInfo, LLT Ty)

It might be safer to create a scalable MMO using Ty rather than forcing it into fixed.

Ditto for the other instructions that are built with LargeSplitSize and SmallSplitSize.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Do you mean we should use LLT LargeSplitSize, SmallSplitSize instead of uint64_t?

Copy link
Contributor

Choose a reason for hiding this comment

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

I think so. Otherwise we risk loosing out on the fact that the size is actually scalable.

@jiahanxie353 jiahanxie353 force-pushed the vload-legalize branch 2 times, most recently from e86058f to 12b0386 Compare July 30, 2024 16:11
@@ -3690,7 +3690,7 @@ LegalizerHelper::LegalizeResult LegalizerHelper::lowerLoad(GAnyLoad &LoadMI) {

if (!isPowerOf2_32(MemSizeInBits)) {
Copy link
Collaborator

Choose a reason for hiding this comment

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

Doesn't this line need to be be updated if MemSizeInBits is a scalable size?

Copy link
Contributor Author

@jiahanxie353 jiahanxie353 Jul 31, 2024

Choose a reason for hiding this comment

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

Indeed, and I'm still pondering how to do it.
Don't have a super clear idea about how to do it. Partly because all of our test cases are captured in the custom branch so they will enter this lower branch (to be honest, I can just revert all my current changes in this file and support it later when there's an actual need); also partly because I don't quite understand why they require this load splitting into power of 2 sized loads.

Copy link
Collaborator

Choose a reason for hiding this comment

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

Ok let's drop the changes if they aren't tested.

@jiahanxie353
Copy link
Contributor Author

Ok let's drop the changes if they aren't tested.

Just dropped it.

Marking it ready for the final review

Copy link
Contributor

@michaelmaitland michaelmaitland left a comment

Choose a reason for hiding this comment

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

LGTM

Copy link
Collaborator

@topperc topperc left a comment

Choose a reason for hiding this comment

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

LGTM

@jiahanxie353 jiahanxie353 merged commit a0d8fa5 into llvm:main Jul 31, 2024
5 of 7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants