[llvm] [LegalizeTypes] Don't pad cttz_elts_zero_poison with ones when widening op (PR #206705)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-aarch64

Author: Luke Lau (lukel97)

<details>
<summary>Changes</summary>

We only need to pad the widened lanes with zeros to handle the case for an all zeroes input. But for cttz_elts_zero_poison, this is already poison.

The RISC-V scalable vector test can't be precomitted because it crashes otherwise trying to lower a get_active_lane_mask from getMaskFromElementCount.

Also while we're here, switch to using TLI.getTypeToTransformTo to be consistent with other `WidenVecOp_*` implementations.

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


3 Files Affected:

- (modified) llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp (+11-5) 
- (modified) llvm/test/CodeGen/AArch64/intrinsic-cttz-elts-sve.ll (+2-9) 
- (modified) llvm/test/CodeGen/RISCV/rvv/cttz-elts.ll (+36) 


``````````diff
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
index bdb8cddda4909..9f7ac8db9c5e4 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
@@ -8730,12 +8730,18 @@ SDValue DAGTypeLegalizer::WidenVecOp_VSELECT(SDNode *N) {
 SDValue DAGTypeLegalizer::WidenVecOp_CttzElements(SDNode *N) {
   SDLoc DL(N);
   SDValue Source = N->getOperand(0);
-  EVT WideVT = GetWidenedVector(Source).getValueType();
+  EVT WideVT =
+      TLI.getTypeToTransformTo(*DAG.getContext(), Source.getValueType());
 
-  // Pad the widened portion with all-ones so the extra lanes appear as
-  // active (non-zero) elements and do not contribute trailing zeros.
-  SDValue AllOnes = DAG.getAllOnesConstant(DL, WideVT);
-  SDValue WideSource = DAG.getInsertSubvector(DL, AllOnes, Source, 0);
+  SDValue WideSource;
+  if (N->getOpcode() == ISD::CTTZ_ELTS_ZERO_POISON)
+    WideSource = GetWidenedVector(Source);
+  else {
+    // Pad the widened portion with all-ones so the extra lanes appear as
+    // active (non-zero) elements and do not contribute trailing zeros.
+    SDValue AllOnes = DAG.getAllOnesConstant(DL, WideVT);
+    WideSource = DAG.getInsertSubvector(DL, AllOnes, Source, 0);
+  }
 
   return DAG.getNode(N->getOpcode(), DL, N->getValueType(0), WideSource,
                      N->getFlags());
diff --git a/llvm/test/CodeGen/AArch64/intrinsic-cttz-elts-sve.ll b/llvm/test/CodeGen/AArch64/intrinsic-cttz-elts-sve.ll
index fa3673feaee46..4a74a2587b39c 100644
--- a/llvm/test/CodeGen/AArch64/intrinsic-cttz-elts-sve.ll
+++ b/llvm/test/CodeGen/AArch64/intrinsic-cttz-elts-sve.ll
@@ -142,15 +142,8 @@ define i32 @ctz_nxv3i1(<vscale x 3 x i1> %a) {
 define i32 @ctz_nxv3i1_poison(<vscale x 3 x i1> %a) {
 ; CHECK-LABEL: ctz_nxv3i1_poison:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    rdvl x8, #1
-; CHECK-NEXT:    mov w9, #3 // =0x3
-; CHECK-NEXT:    ptrue p2.s
-; CHECK-NEXT:    lsr x8, x8, #4
-; CHECK-NEXT:    mul x8, x8, x9
-; CHECK-NEXT:    whilelo p1.s, xzr, x8
-; CHECK-NEXT:    not p1.b, p2/z, p1.b
-; CHECK-NEXT:    mov p0.b, p1/m, p1.b
-; CHECK-NEXT:    brkb p0.b, p2/z, p0.b
+; CHECK-NEXT:    ptrue p1.s
+; CHECK-NEXT:    brkb p0.b, p1/z, p0.b
 ; CHECK-NEXT:    cntp x0, p0, p0.s
 ; CHECK-NEXT:    ret
   %res = call i32 @llvm.experimental.cttz.elts.i32.nxv3i1(<vscale x 3 x i1> %a, i1 1)
diff --git a/llvm/test/CodeGen/RISCV/rvv/cttz-elts.ll b/llvm/test/CodeGen/RISCV/rvv/cttz-elts.ll
index e761f14b0d4e7..99bf2360a0d44 100644
--- a/llvm/test/CodeGen/RISCV/rvv/cttz-elts.ll
+++ b/llvm/test/CodeGen/RISCV/rvv/cttz-elts.ll
@@ -261,4 +261,40 @@ define i32 @ctz_v3i32(<3 x i32> %a) {
   ret i32 %res
 }
 
+define i32 @ctz_v3i32_poison(<3 x i32> %a) {
+; RV32-LABEL: ctz_v3i32_poison:
+; RV32:       # %bb.0:
+; RV32-NEXT:    vsetivli zero, 4, e32, m1, ta, ma
+; RV32-NEXT:    vmsne.vi v8, v8, 0
+; RV32-NEXT:    vfirst.m a0, v8
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: ctz_v3i32_poison:
+; RV64:       # %bb.0:
+; RV64-NEXT:    vsetivli zero, 4, e32, m1, ta, ma
+; RV64-NEXT:    vmsne.vi v8, v8, 0
+; RV64-NEXT:    vfirst.m a0, v8
+; RV64-NEXT:    ret
+  %res = call i32 @llvm.experimental.cttz.elts(<3 x i32> %a, i1 1)
+  ret i32 %res
+}
+
+define i32 @ctz_nxv3i32_poison(<vscale x 3 x i32> %a) {
+; RV32-LABEL: ctz_nxv3i32_poison:
+; RV32:       # %bb.0:
+; RV32-NEXT:    vsetvli a0, zero, e32, m2, ta, ma
+; RV32-NEXT:    vmsne.vi v10, v8, 0
+; RV32-NEXT:    vfirst.m a0, v10
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: ctz_nxv3i32_poison:
+; RV64:       # %bb.0:
+; RV64-NEXT:    vsetvli a0, zero, e32, m2, ta, ma
+; RV64-NEXT:    vmsne.vi v10, v8, 0
+; RV64-NEXT:    vfirst.m a0, v10
+; RV64-NEXT:    ret
+  %res = call i32 @llvm.experimental.cttz.elts(<vscale x 3 x i32> %a, i1 1)
+  ret i32 %res
+}
+
 attributes #0 = { vscale_range(2,1024) }

``````````

</details>


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