[llvm] [DAGCombiner] Handle extending EXTRACT_VECTOR_ELTs in calculateByteProvider (PR #83963)

[Luke Lau via llvm-commits]

https://github.com/lukel97 updated https://github.com/llvm/llvm-project/pull/83963

>From e34658bef9ed282e3ece3c4307b3a73deab71367 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Tue, 5 Mar 2024 14:17:50 +0800
Subject: [PATCH 1/2] [DAGCombiner] Handle extending EXTRACT_VECTOR_ELTs in
 calculateByteProvider

An EXTRACT_VECTOR_ELT can extend the element to the width of its result type,
leaving the high bits undefined. Previously if we attempted to query the bytes
in these high bits we would recurse and hit an assertion. This fixes it by
bailing if the index is outside of the vector element size.

I think the assertion Index < ByteWidth may still be incorrect, since ByteWidth
is calculated from Op.getValueSizeInBits(). I believe this should be
Op.getScalarValueSizeInBits() whenever VectorIndex is set since we're querying
the element now, not the vector. But I couldn't think of a test case to trigger
it. It can be addressed in a follow-up patch.

Fixes #83920
---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp |  4 ++++
 llvm/test/CodeGen/RISCV/rvv/pr83920.ll        | 17 +++++++++++++++++
 2 files changed, 21 insertions(+)
 create mode 100644 llvm/test/CodeGen/RISCV/rvv/pr83920.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 33ada3655dc731..f160dde25d26bf 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -8627,6 +8627,10 @@ calculateByteProvider(SDValue Op, unsigned Index, unsigned Depth,
     if (NarrowBitWidth % 8 != 0)
       return std::nullopt;
     uint64_t NarrowByteWidth = NarrowBitWidth / 8;
+    // EXTRACT_VECTOR_ELT can extend the element type to the width of the return
+    // type, leaving the high bits undefined.
+    if (Index >= NarrowByteWidth)
+      return std::nullopt;
 
     // Check to see if the position of the element in the vector corresponds
     // with the byte we are trying to provide for. In the case of a vector of
diff --git a/llvm/test/CodeGen/RISCV/rvv/pr83920.ll b/llvm/test/CodeGen/RISCV/rvv/pr83920.ll
new file mode 100644
index 00000000000000..058417629c9c4e
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/rvv/pr83920.ll
@@ -0,0 +1,17 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
+; RUN: llc < %s -mtriple=riscv64 -mattr=+v -verify-machineinstrs | FileCheck %s
+
+define i8 @or_load_combine(ptr %p) {
+; CHECK-LABEL: or_load_combine:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    vsetivli zero, 2, e8, mf8, ta, ma
+; CHECK-NEXT:    vle8.v v8, (a0)
+; CHECK-NEXT:    vmv.x.s a0, v8
+; CHECK-NEXT:    ori a0, a0, 1
+; CHECK-NEXT:    ret
+entry:
+  %0 = load <2 x i8>, ptr %p
+  %1 = extractelement <2 x i8> %0, i64 0
+  %2 = or i8 %1, 1
+  ret i8 %2
+}

>From 3cf1b40c18aa8a960df0550c0e026e14bd33ebfe Mon Sep 17 00:00:00 2001
From: Luke Lau <luke at igalia.com>
Date: Tue, 5 Mar 2024 18:30:33 +0800
Subject: [PATCH 2/2] Name variables, drop -verify-machineinstrs

---
 llvm/test/CodeGen/RISCV/rvv/pr83920.ll | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/llvm/test/CodeGen/RISCV/rvv/pr83920.ll b/llvm/test/CodeGen/RISCV/rvv/pr83920.ll
index 058417629c9c4e..319f0bd9ef1cd7 100644
--- a/llvm/test/CodeGen/RISCV/rvv/pr83920.ll
+++ b/llvm/test/CodeGen/RISCV/rvv/pr83920.ll
@@ -1,17 +1,16 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
-; RUN: llc < %s -mtriple=riscv64 -mattr=+v -verify-machineinstrs | FileCheck %s
+; RUN: llc < %s -mtriple=riscv64 -mattr=+v | FileCheck %s
 
 define i8 @or_load_combine(ptr %p) {
 ; CHECK-LABEL: or_load_combine:
-; CHECK:       # %bb.0: # %entry
+; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetivli zero, 2, e8, mf8, ta, ma
 ; CHECK-NEXT:    vle8.v v8, (a0)
 ; CHECK-NEXT:    vmv.x.s a0, v8
 ; CHECK-NEXT:    ori a0, a0, 1
 ; CHECK-NEXT:    ret
-entry:
-  %0 = load <2 x i8>, ptr %p
-  %1 = extractelement <2 x i8> %0, i64 0
-  %2 = or i8 %1, 1
-  ret i8 %2
+  %load = load <2 x i8>, ptr %p
+  %extract = extractelement <2 x i8> %load, i64 0
+  %or = or i8 %extract, 1
+  ret i8 %or
 }