[GlobalIsel] Combine extract vector element #90339

tschuett · 2024-04-27T11:11:04Z

look through shuffle vectors

llvmbot · 2024-04-27T11:11:33Z

@llvm/pr-subscribers-llvm-globalisel

@llvm/pr-subscribers-backend-aarch64

Author: Thorsten Schütt (tschuett)

Changes

look through shuffle vectors

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

6 Files Affected:

(modified) llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h (+5)
(modified) llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h (+12)
(modified) llvm/include/llvm/Target/GlobalISel/Combine.td (+8)
(modified) llvm/lib/CodeGen/GlobalISel/CombinerHelperVectorOps.cpp (+107)
(modified) llvm/test/CodeGen/AArch64/GlobalISel/combine-extract-vec-elt.mir (+81)
(modified) llvm/test/CodeGen/AArch64/extract-vector-elt.ll (+4-15)

diff --git a/llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h b/llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h
index 76e8d1166ae0cd..4f1c9642e117c2 100644
--- a/llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h
+++ b/llvm/include/llvm/CodeGen/GlobalISel/CombinerHelper.h
@@ -848,6 +848,11 @@ class CombinerHelper {
   bool matchExtractVectorElementWithBuildVectorTrunc(const MachineOperand &MO,
                                                      BuildFnTy &MatchInfo);
 
+  /// Combine extract vector element with a shuffle vector on the vector
+  /// register.
+  bool matchExtractVectorElementWithShuffleVector(const MachineOperand &MO,
+                                                  BuildFnTy &MatchInfo);
+
   /// Combine extract vector element with a insert vector element on the vector
   /// register and different indices.
   bool matchExtractVectorElementWithDifferentIndices(const MachineOperand &MO,
diff --git a/llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h b/llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h
index 25e47114e4a39a..47ad0c3d6430e8 100644
--- a/llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h
+++ b/llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h
@@ -294,6 +294,18 @@ class GBuildVectorTrunc : public GMergeLikeInstr {
   }
 };
 
+/// Respresents a G_SHUFFLE_VECTOR.
+class GShuffleVector : public GenericMachineInstr {
+public:
+  Register getSrc1Reg() const { return getOperand(1).getReg(); }
+  Register getSrc2Reg() const { return getOperand(2).getReg(); }
+  llvm::ArrayRef<int> getMask() const { return getOperand(3).getShuffleMask(); }
+
+  static bool classof(const MachineInstr *MI) {
+    return MI->getOpcode() == TargetOpcode::G_SHUFFLE_VECTOR;
+  }
+};
+
 /// Represents a G_PTR_ADD.
 class GPtrAdd : public GenericMachineInstr {
 public:
diff --git a/llvm/include/llvm/Target/GlobalISel/Combine.td b/llvm/include/llvm/Target/GlobalISel/Combine.td
index dbbb3abaa83047..55c17014e29808 100644
--- a/llvm/include/llvm/Target/GlobalISel/Combine.td
+++ b/llvm/include/llvm/Target/GlobalISel/Combine.td
@@ -1505,6 +1505,13 @@ def extract_vector_element_freeze : GICombineRule<
    [{ return Helper.matchExtractVectorElementWithFreeze(${root}, ${matchinfo}); }]),
    (apply [{ Helper.applyBuildFnMO(${root}, ${matchinfo}); }])>;
 
+def extract_vector_element_shuffle_vector : GICombineRule<
+   (defs root:$root, build_fn_matchinfo:$matchinfo),
+   (match (G_SHUFFLE_VECTOR $src, $src1, $src2, $mask),
+          (G_EXTRACT_VECTOR_ELT $root, $src, $idx),
+   [{ return Helper.matchExtractVectorElementWithShuffleVector(${root}, ${matchinfo}); }]),
+   (apply [{ Helper.applyBuildFnMO(${root}, ${matchinfo}); }])>;
+
 // Combines concat operations
 def concat_matchinfo : GIDefMatchData<"SmallVector<Register>">;
 def combine_concat_vector : GICombineRule<
@@ -1582,6 +1589,7 @@ extract_vector_element_build_vector_trunc6,
 extract_vector_element_build_vector_trunc7,
 extract_vector_element_build_vector_trunc8,
 extract_vector_element_freeze,
+extract_vector_element_shuffle_vector,
 insert_vector_element_extract_vector_element
 ]>;
 
diff --git a/llvm/lib/CodeGen/GlobalISel/CombinerHelperVectorOps.cpp b/llvm/lib/CodeGen/GlobalISel/CombinerHelperVectorOps.cpp
index fb33801a3a33fc..e1ba49ec2f0fa9 100644
--- a/llvm/lib/CodeGen/GlobalISel/CombinerHelperVectorOps.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/CombinerHelperVectorOps.cpp
@@ -325,6 +325,113 @@ bool CombinerHelper::matchExtractVectorElementWithBuildVectorTrunc(
   return true;
 }
 
+bool CombinerHelper::matchExtractVectorElementWithShuffleVector(
+    const MachineOperand &MO, BuildFnTy &MatchInfo) {
+  GExtractVectorElement *Extract =
+      cast<GExtractVectorElement>(getDefIgnoringCopies(MO.getReg(), MRI));
+
+  //
+  //  %zero:_(s64) = G_CONSTANT i64 0
+  //  %sv:_(<4 x s32>) = G_SHUFFLE_SHUFFLE %arg1(<4 x s32>), %arg2(<4 x s32>),
+  //                     shufflemask(0, 0, 0, 0)
+  //  %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %zero(s64)
+  //
+  //  -->
+  //
+  //  %zero1:_(s64) = G_CONSTANT i64 0
+  //  %extract:_(s32) = G_EXTRACT_VECTOR_ELT %arg1(<4 x s32>), %zero1(s64)
+  //
+  //
+  //
+  //
+  //  %three:_(s64) = G_CONSTANT i64 3
+  //  %sv:_(<4 x s32>) = G_SHUFFLE_SHUFFLE %arg1(<4 x s32>), %arg2(<4 x s32>),
+  //                     shufflemask(0, 0, 0, -1)
+  //  %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %three(s64)
+  //
+  //  -->
+  //
+  //  %extract:_(s32) = G_IMPLICIT_DEF
+  //
+  //
+  //
+  //
+  //
+  //  %sv:_(<4 x s32>) = G_SHUFFLE_SHUFFLE %arg1(<4 x s32>), %arg2(<4 x s32>),
+  //                     shufflemask(0, 0, 0, -1)
+  //  %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %opaque(s64)
+  //
+  //  -->
+  //
+  //  %sv:_(<4 x s32>) = G_SHUFFLE_SHUFFLE %arg1(<4 x s32>), %arg2(<4 x s32>),
+  //                     shufflemask(0, 0, 0, -1)
+  //  %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %opaque(s64)
+  //
+  //
+
+  // We try to get the value of the Index register.
+  std::optional<ValueAndVReg> MaybeIndex =
+      getIConstantVRegValWithLookThrough(Extract->getIndexReg(), MRI);
+  if (!MaybeIndex)
+    return false;
+
+  GShuffleVector *Shuffle =
+      cast<GShuffleVector>(getDefIgnoringCopies(Extract->getVectorReg(), MRI));
+
+  llvm::ArrayRef<int> Mask = Shuffle->getMask();
+
+  unsigned Offset = MaybeIndex->Value.getZExtValue();
+  int SrcIdx = Mask[Offset];
+
+  LLT Src1Type = MRI.getType(Shuffle->getSrc1Reg());
+  // At the IR level a <1 x ty> shuffle  vector is valid, but we want to extract
+  // from a vector.
+  assert(Src1Type.isVector() && "expected to extract from a vector");
+  unsigned LHSWidth = Src1Type.isVector() ? Src1Type.getNumElements() : 1;
+
+  // Note that there is no one use check.
+  Register Dst = Extract->getReg(0);
+  LLT DstTy = MRI.getType(Dst);
+
+  if (SrcIdx < 0 &&
+      isLegalOrBeforeLegalizer({TargetOpcode::G_IMPLICIT_DEF, {DstTy}})) {
+    MatchInfo = [=](MachineIRBuilder &B) { B.buildUndef(Dst); };
+    return true;
+  }
+
+  // If the legality check failed, then we still have to abort.
+  if (SrcIdx < 0)
+    return false;
+
+  Register NewVector;
+
+  // We check in which vector and at what offset to look through.
+  if (SrcIdx < (int)LHSWidth) {
+    NewVector = Shuffle->getSrc1Reg();
+    // SrcIdx unchanged
+  } else { // SrcIdx >= LHSWidth
+    NewVector = Shuffle->getSrc2Reg();
+    SrcIdx -= LHSWidth;
+  }
+
+  LLT IdxTy = MRI.getType(Extract->getIndexReg());
+  LLT NewVectorTy = MRI.getType(NewVector);
+
+  // We check the legality of the look through.
+  if (!isLegalOrBeforeLegalizer(
+          {TargetOpcode::G_EXTRACT_VECTOR_ELT, {DstTy, NewVectorTy, IdxTy}}) ||
+      !isConstantLegalOrBeforeLegalizer({IdxTy}))
+    return false;
+
+  // We look through the shuffle vector.
+  MatchInfo = [=](MachineIRBuilder &B) {
+    auto Idx = B.buildConstant(IdxTy, SrcIdx);
+    B.buildExtractVectorElement(Dst, NewVector, Idx);
+  };
+
+  return true;
+}
+
 bool CombinerHelper::matchInsertVectorElementOOB(MachineInstr &MI,
                                                  BuildFnTy &MatchInfo) {
   GInsertVectorElement *Insert = cast<GInsertVectorElement>(&MI);
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/combine-extract-vec-elt.mir b/llvm/test/CodeGen/AArch64/GlobalISel/combine-extract-vec-elt.mir
index 587d53c300f83f..3b7515b407e37f 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/combine-extract-vec-elt.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/combine-extract-vec-elt.mir
@@ -571,3 +571,84 @@ body:             |
     RET_ReallyLR implicit $x0
 ...
 ---
+name:            extract_from_build_vector_shuffle_vector_undef
+alignment:       4
+liveins:
+  - { reg: '$x0' }
+  - { reg: '$x1' }
+frameInfo:
+  maxAlignment:    1
+body:             |
+  bb.1:
+    liveins: $x0, $x1
+    ; CHECK-LABEL: name: extract_from_build_vector_shuffle_vector_undef
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %extract:_(s32) = G_IMPLICIT_DEF
+    ; CHECK-NEXT: $w0 = COPY %extract(s32)
+    ; CHECK-NEXT: RET_ReallyLR implicit $x0
+    %arg1:_(<4 x s32>) = COPY $q0
+    %arg2:_(<4 x s32>) = COPY $q1
+    %idx:_(s64) = G_CONSTANT i64 0
+    %sv:_(<4 x s32>) = G_SHUFFLE_VECTOR %arg1(<4 x s32>), %arg2(<4 x s32>), shufflemask(-1, 0, 0, 0)
+    %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %idx(s64)
+    $w0 = COPY %extract(s32)
+    RET_ReallyLR implicit $x0
+...
+---
+name:            extract_from_build_vector_shuffle_vector_opaque
+alignment:       4
+liveins:
+  - { reg: '$x0' }
+  - { reg: '$x1' }
+frameInfo:
+  maxAlignment:    1
+body:             |
+  bb.1:
+    liveins: $x0, $x1
+    ; CHECK-LABEL: name: extract_from_build_vector_shuffle_vector_opaque
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %arg1:_(<4 x s32>) = COPY $q0
+    ; CHECK-NEXT: %arg2:_(<4 x s32>) = COPY $q1
+    ; CHECK-NEXT: %idx:_(s64) = COPY $x1
+    ; CHECK-NEXT: %sv:_(<4 x s32>) = G_SHUFFLE_VECTOR %arg1(<4 x s32>), %arg2, shufflemask(undef, 0, 0, 0)
+    ; CHECK-NEXT: %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %idx(s64)
+    ; CHECK-NEXT: $w0 = COPY %extract(s32)
+    ; CHECK-NEXT: RET_ReallyLR implicit $x0
+    %arg1:_(<4 x s32>) = COPY $q0
+    %arg2:_(<4 x s32>) = COPY $q1
+    %idx:_(s64) = COPY $x1
+    %sv:_(<4 x s32>) = G_SHUFFLE_VECTOR %arg1(<4 x s32>), %arg2(<4 x s32>), shufflemask(-1, 0, 0, 0)
+    %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %idx(s64)
+    $w0 = COPY %extract(s32)
+    RET_ReallyLR implicit $x0
+...
+---
+name:            extract_from_build_vector_shuffle_vector_const
+alignment:       4
+liveins:
+  - { reg: '$x0' }
+  - { reg: '$x1' }
+frameInfo:
+  maxAlignment:    1
+body:             |
+  bb.1:
+    liveins: $x0, $x1
+    ; CHECK-LABEL: name: extract_from_build_vector_shuffle_vector_const
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %arg1:_(<4 x s32>) = COPY $q0
+    ; CHECK-NEXT: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 3
+    ; CHECK-NEXT: %extract:_(s32) = G_EXTRACT_VECTOR_ELT %arg1(<4 x s32>), [[C]](s64)
+    ; CHECK-NEXT: $w0 = COPY %extract(s32)
+    ; CHECK-NEXT: RET_ReallyLR implicit $x0
+    %arg1:_(<4 x s32>) = COPY $q0
+    %arg2:_(<4 x s32>) = COPY $q1
+    %idx:_(s64) = G_CONSTANT i64 0
+    %sv:_(<4 x s32>) = G_SHUFFLE_VECTOR %arg1(<4 x s32>), %arg2(<4 x s32>), shufflemask(3, 0, 0, 0)
+    %extract:_(s32) = G_EXTRACT_VECTOR_ELT %sv(<4 x s32>), %idx(s64)
+    $w0 = COPY %extract(s32)
+    RET_ReallyLR implicit $x0
+...
+---
diff --git a/llvm/test/CodeGen/AArch64/extract-vector-elt.ll b/llvm/test/CodeGen/AArch64/extract-vector-elt.ll
index 504222e0036e22..0481d997d24faf 100644
--- a/llvm/test/CodeGen/AArch64/extract-vector-elt.ll
+++ b/llvm/test/CodeGen/AArch64/extract-vector-elt.ll
@@ -938,21 +938,10 @@ entry:
 }
 
 define i32 @extract_v4i32_shuffle_const(<4 x i32> %a, <4 x i32> %b, i32 %c) {
-; CHECK-SD-LABEL: extract_v4i32_shuffle_const:
-; CHECK-SD:       // %bb.0: // %entry
-; CHECK-SD-NEXT:    fmov w0, s1
-; CHECK-SD-NEXT:    ret
-;
-; CHECK-GI-LABEL: extract_v4i32_shuffle_const:
-; CHECK-GI:       // %bb.0: // %entry
-; CHECK-GI-NEXT:    adrp x8, .LCPI36_0
-; CHECK-GI-NEXT:    // kill: def $q0 killed $q0 killed $q0_q1 def $q0_q1
-; CHECK-GI-NEXT:    ldr q2, [x8, :lo12:.LCPI36_0]
-; CHECK-GI-NEXT:    // kill: def $q1 killed $q1 killed $q0_q1 def $q0_q1
-; CHECK-GI-NEXT:    tbl v0.16b, { v0.16b, v1.16b }, v2.16b
-; CHECK-GI-NEXT:    mov s0, v0.s[2]
-; CHECK-GI-NEXT:    fmov w0, s0
-; CHECK-GI-NEXT:    ret
+; CHECK-LABEL: extract_v4i32_shuffle_const:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    fmov w0, s1
+; CHECK-NEXT:    ret
 entry:
   %vector = shufflevector <4 x i32> %a, <4 x i32> %b, <4 x i32> <i32 0, i32 2, i32 4, i32 3>
   %d = extractelement <4 x i32> %vector, i32 2

llvm/lib/CodeGen/GlobalISel/CombinerHelperVectorOps.cpp

arsenm · 2024-05-06T08:36:11Z

llvm/lib/CodeGen/GlobalISel/CombinerHelperVectorOps.cpp

+  GShuffleVector *Shuffle =
+      cast<GShuffleVector>(getDefIgnoringCopies(Extract->getVectorReg(), MRI));


It's unfortunate you need to re-find the def instruction when it was already matched

Yeah, but the cast and get def is safer than I KNOW WHAT I AM DOING.

llvm/test/CodeGen/AArch64/GlobalISel/combine-extract-vec-elt.mir

llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h

[GlobalIsel] Combine extract vector element

5a911b1

look through shuffle vectors

tschuett requested review from aemerson and arsenm April 27, 2024 11:11

llvmbot added backend:AArch64 llvm:globalisel labels Apr 27, 2024

arsenm reviewed May 6, 2024

View reviewed changes

address review comments

b63319a

madhur13490 reviewed May 6, 2024

View reviewed changes

llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h Outdated Show resolved Hide resolved

Thorsten Schütt added 2 commits May 6, 2024 16:36

fix typo

310c82b

retrigger buildkite

5797a33

aemerson approved these changes May 6, 2024

View reviewed changes

tschuett merged commit b42f553 into llvm:main May 7, 2024
4 checks passed

tschuett deleted the gisel-extract-vector-element-sv branch May 7, 2024 05:13

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[GlobalIsel] Combine extract vector element #90339

[GlobalIsel] Combine extract vector element #90339

Uh oh!

tschuett commented Apr 27, 2024

Uh oh!

llvmbot commented Apr 27, 2024 •

edited

Loading

Uh oh!

Uh oh!

arsenm May 6, 2024

Uh oh!

tschuett May 6, 2024

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

		GShuffleVector *Shuffle =
		cast<GShuffleVector>(getDefIgnoringCopies(Extract->getVectorReg(), MRI));

[GlobalIsel] Combine extract vector element #90339

[GlobalIsel] Combine extract vector element #90339

Uh oh!

Conversation

tschuett commented Apr 27, 2024

Uh oh!

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

Uh oh!

Uh oh!

Uh oh!

arsenm May 6, 2024

Choose a reason for hiding this comment

Uh oh!

tschuett May 6, 2024

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

llvmbot commented Apr 27, 2024 •

edited

Loading