[llvm] [RISCV][llvm] Correct shamt in P extension EXTRACT_VECTOR_ELT lowering (PR #169823)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

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

Author: Brandon Wu (4vtomat)

<details>
<summary>Changes</summary>

During operation legalization, element type should have been turn into
XLenVT which makes the SHL a no-op. We need to use exact vector element
type instead.


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


3 Files Affected:

- (modified) llvm/lib/Target/RISCV/RISCVISelLowering.cpp (+1-1) 
- (modified) llvm/test/CodeGen/RISCV/rvp-ext-rv32.ll (+32) 
- (modified) llvm/test/CodeGen/RISCV/rvp-ext-rv64.ll (+12) 


``````````diff
diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index be53f51afe79f..0d9b8bddbcc8f 100644
--- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -10699,7 +10699,7 @@ SDValue RISCVTargetLowering::lowerEXTRACT_VECTOR_ELT(SDValue Op,
         VecVT != MVT::v4i8 && VecVT != MVT::v2i32)
       return SDValue();
     SDValue Extracted = DAG.getBitcast(XLenVT, Vec);
-    unsigned ElemWidth = EltVT.getSizeInBits();
+    unsigned ElemWidth = VecVT.getVectorElementType().getSizeInBits();
     SDValue Shamt = DAG.getNode(ISD::MUL, DL, XLenVT, Idx,
                                 DAG.getConstant(ElemWidth, DL, XLenVT));
     return DAG.getNode(ISD::SRL, DL, XLenVT, Extracted, Shamt);
diff --git a/llvm/test/CodeGen/RISCV/rvp-ext-rv32.ll b/llvm/test/CodeGen/RISCV/rvp-ext-rv32.ll
index d4ea9e6c3def0..f803f6aa09652 100644
--- a/llvm/test/CodeGen/RISCV/rvp-ext-rv32.ll
+++ b/llvm/test/CodeGen/RISCV/rvp-ext-rv32.ll
@@ -484,6 +484,25 @@ define void @test_extract_vector_16(ptr %ret_ptr, ptr %a_ptr) {
   ret void
 }
 
+define void @test_extract_vector_16_elem1(ptr %ret_ptr, ptr %a_ptr) {
+; CHECK-RV32-LABEL: test_extract_vector_16_elem1:
+; CHECK-RV32:       # %bb.0:
+; CHECK-RV32-NEXT:    lhu a1, 2(a1)
+; CHECK-RV32-NEXT:    sh a1, 0(a0)
+; CHECK-RV32-NEXT:    ret
+;
+; CHECK-RV64-LABEL: test_extract_vector_16_elem1:
+; CHECK-RV64:       # %bb.0:
+; CHECK-RV64-NEXT:    lw a1, 0(a1)
+; CHECK-RV64-NEXT:    srli a1, a1, 16
+; CHECK-RV64-NEXT:    sh a1, 0(a0)
+; CHECK-RV64-NEXT:    ret
+  %a = load <2 x i16>, ptr %a_ptr
+  %extracted = extractelement <2 x i16> %a, i32 1
+  store i16 %extracted, ptr %ret_ptr
+  ret void
+}
+
 define void @test_extract_vector_8(ptr %ret_ptr, ptr %a_ptr) {
 ; CHECK-LABEL: test_extract_vector_8:
 ; CHECK:       # %bb.0:
@@ -496,6 +515,19 @@ define void @test_extract_vector_8(ptr %ret_ptr, ptr %a_ptr) {
   ret void
 }
 
+define void @test_extract_vector_8_elem1(ptr %ret_ptr, ptr %a_ptr) {
+; CHECK-LABEL: test_extract_vector_8_elem1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    lw a1, 0(a1)
+; CHECK-NEXT:    srli a1, a1, 8
+; CHECK-NEXT:    sb a1, 0(a0)
+; CHECK-NEXT:    ret
+  %a = load <4 x i8>, ptr %a_ptr
+  %extracted = extractelement <4 x i8> %a, i32 1
+  store i8 %extracted, ptr %ret_ptr
+  ret void
+}
+
 ; Test for splat
 define void @test_non_const_splat_i8(ptr %ret_ptr, ptr %a_ptr, i8 %elt) {
 ; CHECK-LABEL: test_non_const_splat_i8:
diff --git a/llvm/test/CodeGen/RISCV/rvp-ext-rv64.ll b/llvm/test/CodeGen/RISCV/rvp-ext-rv64.ll
index b39b807d43154..9b021df8dd452 100644
--- a/llvm/test/CodeGen/RISCV/rvp-ext-rv64.ll
+++ b/llvm/test/CodeGen/RISCV/rvp-ext-rv64.ll
@@ -495,6 +495,18 @@ define void @test_extract_vector_32(ptr %ret_ptr, ptr %a_ptr) {
   ret void
 }
 
+define void @test_extract_vector_32_elem1(ptr %ret_ptr, ptr %a_ptr) {
+; CHECK-LABEL: test_extract_vector_32_elem1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    lw a1, 4(a1)
+; CHECK-NEXT:    sw a1, 0(a0)
+; CHECK-NEXT:    ret
+  %a = load <2 x i32>, ptr %a_ptr
+  %extracted = extractelement <2 x i32> %a, i32 1
+  store i32 %extracted, ptr %ret_ptr
+  ret void
+}
+
 ; Test basic add/sub operations for v2i32 (RV64 only)
 define void @test_padd_w(ptr %ret_ptr, ptr %a_ptr, ptr %b_ptr) {
 ; CHECK-LABEL: test_padd_w:

``````````

</details>


https://github.com/llvm/llvm-project/pull/169823