release/19.x: [DAG] Allow AssertZExt to scalarize. (#122463)

[llvmbot]

Backport ab9a80a

Requested by: @nikic

[llvmbot]

@arsenm What do you think about merging this PR to the release branch?

[llvmbot]

@llvm/pr-subscribers-llvm-selectiondag

@llvm/pr-subscribers-backend-arm

Author: None (llvmbot)

Changes

Backport ab9a80a

Requested by: @nikic

---

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

3 Files Affected:

• (modified) llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h (+1-1)
• (modified) llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp (+6-2)
• (added) llvm/test/CodeGen/ARM/scalarize-assert-zext.ll (+46)

diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
index d4e61c85889012..d74896772bf53b 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
@@ -838,7 +838,7 @@ class LLVM_LIBRARY_VISIBILITY DAGTypeLegalizer {
   SDValue ScalarizeVecRes_BUILD_VECTOR(SDNode *N);
   SDValue ScalarizeVecRes_EXTRACT_SUBVECTOR(SDNode *N);
   SDValue ScalarizeVecRes_FP_ROUND(SDNode *N);
-  SDValue ScalarizeVecRes_ExpOp(SDNode *N);
+  SDValue ScalarizeVecRes_UnaryOpWithExtraInput(SDNode *N);
   SDValue ScalarizeVecRes_INSERT_VECTOR_ELT(SDNode *N);
   SDValue ScalarizeVecRes_LOAD(LoadSDNode *N);
   SDValue ScalarizeVecRes_SCALAR_TO_VECTOR(SDNode *N);
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
index 92b62ccdc27552..ea95aaef8a1e87 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
@@ -58,7 +58,11 @@ void DAGTypeLegalizer::ScalarizeVectorResult(SDNode *N, unsigned ResNo) {
   case ISD::BUILD_VECTOR:      R = ScalarizeVecRes_BUILD_VECTOR(N); break;
   case ISD::EXTRACT_SUBVECTOR: R = ScalarizeVecRes_EXTRACT_SUBVECTOR(N); break;
   case ISD::FP_ROUND:          R = ScalarizeVecRes_FP_ROUND(N); break;
-  case ISD::FPOWI:             R = ScalarizeVecRes_ExpOp(N); break;
+  case ISD::AssertZext:
+  case ISD::AssertSext:
+  case ISD::FPOWI:
+    R = ScalarizeVecRes_UnaryOpWithExtraInput(N);
+    break;
   case ISD::INSERT_VECTOR_ELT: R = ScalarizeVecRes_INSERT_VECTOR_ELT(N); break;
   case ISD::LOAD:           R = ScalarizeVecRes_LOAD(cast<LoadSDNode>(N));break;
   case ISD::SCALAR_TO_VECTOR:  R = ScalarizeVecRes_SCALAR_TO_VECTOR(N); break;
@@ -426,7 +430,7 @@ SDValue DAGTypeLegalizer::ScalarizeVecRes_FP_ROUND(SDNode *N) {
                      N->getOperand(1));
 }
 
-SDValue DAGTypeLegalizer::ScalarizeVecRes_ExpOp(SDNode *N) {
+SDValue DAGTypeLegalizer::ScalarizeVecRes_UnaryOpWithExtraInput(SDNode *N) {
   SDValue Op = GetScalarizedVector(N->getOperand(0));
   return DAG.getNode(N->getOpcode(), SDLoc(N), Op.getValueType(), Op,
                      N->getOperand(1));
diff --git a/llvm/test/CodeGen/ARM/scalarize-assert-zext.ll b/llvm/test/CodeGen/ARM/scalarize-assert-zext.ll
new file mode 100644
index 00000000000000..5638bb4a398803
--- /dev/null
+++ b/llvm/test/CodeGen/ARM/scalarize-assert-zext.ll
@@ -0,0 +1,46 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -mtriple=armv7-unknown-linux-musleabihf -mattr=-neon %s -o - | FileCheck %s
+
+declare fastcc noundef range(i16 0, 256) <4 x i16> @other()
+
+define void @test(ptr %0) #0 {
+; CHECK-LABEL: test:
+; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    .save {r4, lr}
+; CHECK-NEXT:    push {r4, lr}
+; CHECK-NEXT:    mov r4, r0
+; CHECK-NEXT:    bl other
+; CHECK-NEXT:    uxth r3, r3
+; CHECK-NEXT:    uxth r2, r2
+; CHECK-NEXT:    uxth r1, r1
+; CHECK-NEXT:    uxth r0, r0
+; CHECK-NEXT:    strb r3, [r4, #3]
+; CHECK-NEXT:    strb r2, [r4, #2]
+; CHECK-NEXT:    strb r1, [r4, #1]
+; CHECK-NEXT:    strb r0, [r4]
+; CHECK-NEXT:    pop {r4, pc}
+entry:
+  %call = call fastcc <4 x i16> @other()
+  %t = trunc <4 x i16> %call to <4 x i8>
+  store <4 x i8> %t, ptr %0, align 1
+  ret void
+}
+
+define <4 x i16> @test2() #0 {
+; CHECK-LABEL: test2:
+; CHECK:       @ %bb.0: @ %entry
+; CHECK-NEXT:    .save {r11, lr}
+; CHECK-NEXT:    push {r11, lr}
+; CHECK-NEXT:    bl other
+; CHECK-NEXT:    movw r1, #65408
+; CHECK-NEXT:    and r0, r0, r1
+; CHECK-NEXT:    and r2, r2, r1
+; CHECK-NEXT:    mov r1, #0
+; CHECK-NEXT:    mov r3, #0
+; CHECK-NEXT:    pop {r11, pc}
+entry:
+  %call = call fastcc <4 x i16> @other()
+  %a = and <4 x i16> %call, <i16 u0x80, i16 u0x100, i16 u0x80, i16 u0x100>
+  ret <4 x i16> %a
+}
+

[github-actions]

@nikic (or anyone else). If you would like to add a note about this fix in the release notes (completely optional). Please reply to this comment with a one or two sentence description of the fix. When you are done, please add the release:note label to this PR.