[llvm] [DAGCombiner] Fold trunc(build_vector(ext(x), ext(x)) -> build_vector(x,x) (PR #179857)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-selectiondag

Author: Liao Chunyu (ChunyuLiao)

<details>
<summary>Changes</summary>

 The original implementation performed the transformation when isTruncateFree was true:
 truncate(build_vector(x, x)) -> build_vector(truncate(x), truncate(x)).
    
 In some cases, x comes from an ext, try to pre-truncate build_vectors source operands
 when the source operands of build_vectors comes from an ext.
    
 Testcase from: https://gcc.godbolt.org/z/bbxbYK7dh

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


2 Files Affected:

- (modified) llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp (+19-5) 
- (modified) llvm/test/CodeGen/RISCV/rvv/fixed-vectors-vaaddu.ll (+39) 


``````````diff
diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index df69f0870d27a..772dedf5f79cd 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -16634,16 +16634,30 @@ SDValue DAGCombiner::visitTRUNCATE(SDNode *N) {
   // Attempt to pre-truncate BUILD_VECTOR sources.
   if (N0.getOpcode() == ISD::BUILD_VECTOR && !LegalOperations &&
       N0.hasOneUse() &&
-      TLI.isTruncateFree(SrcVT.getScalarType(), VT.getScalarType()) &&
       // Avoid creating illegal types if running after type legalizer.
       (!LegalTypes || TLI.isTypeLegal(VT.getScalarType()))) {
     EVT SVT = VT.getScalarType();
     SmallVector<SDValue, 8> TruncOps;
-    for (const SDValue &Op : N0->op_values()) {
-      SDValue TruncOp = DAG.getNode(ISD::TRUNCATE, DL, SVT, Op);
-      TruncOps.push_back(TruncOp);
+
+    if (TLI.isTruncateFree(SrcVT.getScalarType(), VT.getScalarType())) {
+      for (const SDValue &Op : N0->op_values()) {
+        SDValue TruncOp = DAG.getNode(ISD::TRUNCATE, DL, SVT, Op);
+        TruncOps.push_back(TruncOp);
+      }
+      return DAG.getBuildVector(VT, DL, TruncOps);
+    }
+    // trunc(build_vector(ext(x), ext(x)) -> build_vector(x,x)
+    SDValue SplatVal = DAG.getSplatValue(N0);
+    if (SplatVal) {
+      unsigned Opcode = SplatVal.getOpcode();
+      if ((Opcode == ISD::SIGN_EXTEND || Opcode == ISD::ZERO_EXTEND ||
+           Opcode == ISD::ANY_EXTEND) &&
+          SrcVT.getScalarType() == SplatVal.getValueType()) {
+        for (const SDValue &Op : N0->op_values())
+          TruncOps.push_back(Op);
+        return DAG.getBuildVector(VT, DL, TruncOps);
+      }
     }
-    return DAG.getBuildVector(VT, DL, TruncOps);
   }
 
   // trunc (splat_vector x) -> splat_vector (trunc x)
diff --git a/llvm/test/CodeGen/RISCV/rvv/fixed-vectors-vaaddu.ll b/llvm/test/CodeGen/RISCV/rvv/fixed-vectors-vaaddu.ll
index dc432efbd5c47..755ad8c402628 100644
--- a/llvm/test/CodeGen/RISCV/rvv/fixed-vectors-vaaddu.ll
+++ b/llvm/test/CodeGen/RISCV/rvv/fixed-vectors-vaaddu.ll
@@ -212,6 +212,45 @@ define <8 x i64> @vaaddu_vx_v8i64_floor(<8 x i64> %x, i64 %y) {
   ret <8 x i64> %ret
 }
 
+define <8 x i8> @vaaddu_vx_floor_splat_scalar(<8 x i8> %a, i8 %sc) {
+; CHECK-LABEL: vaaddu_vx_floor_splat_scalar:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    xori a0, a0, 130
+; CHECK-NEXT:    csrwi vxrm, 2
+; CHECK-NEXT:    vsetivli zero, 8, e8, mf2, ta, ma
+; CHECK-NEXT:    vaaddu.vx v8, v8, a0
+; CHECK-NEXT:    ret
+  %1 = xor i8 %sc, -126
+  %c = zext i8 %1 to i16
+  %za = zext <8 x i8> %a to <8 x i16>
+  %yhead = insertelement <8 x i16> poison, i16 %c, i64 0
+  %yzv = shufflevector <8 x i16> %yhead, <8 x i16> poison, <8 x i32> zeroinitializer
+  %add = add nuw nsw <8 x i16> %yzv, %za
+  %div = lshr <8 x i16> %add, splat (i16 1)
+  %ret = trunc nuw <8 x i16> %div to <8 x i8>
+  ret <8 x i8> %ret
+}
+
+define <8 x i8> @vaaddu_vx_ceil_splat_scalar(<8 x i8> %x, i8 %y) {
+; CHECK-LABEL: vaaddu_vx_ceil_splat_scalar:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    xori a0, a0, 130
+; CHECK-NEXT:    csrwi vxrm, 0
+; CHECK-NEXT:    vsetivli zero, 8, e8, mf2, ta, ma
+; CHECK-NEXT:    vaaddu.vx v8, v8, a0
+; CHECK-NEXT:    ret
+  %1 = xor i8 %y, -126
+  %c = zext i8 %1 to i16
+  %xzv = zext <8 x i8> %x to <8 x i16>
+  %yhead = insertelement <8 x i16> poison, i16 %c, i32 0
+  %yzv = shufflevector <8 x i16> %yhead, <8 x i16> poison, <8 x i32> zeroinitializer
+  %add = add nuw nsw <8 x i16> %xzv, %yzv
+  %add1 = add nuw nsw <8 x i16> %add, <i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1, i16 1>
+  %div = lshr <8 x i16> %add1, splat (i16 1)
+  %ret = trunc <8 x i16> %div to <8 x i8>
+  ret <8 x i8> %ret
+}
+
 define <8 x i8> @vaaddu_vv_v8i8_ceil(<8 x i8> %x, <8 x i8> %y) {
 ; CHECK-LABEL: vaaddu_vv_v8i8_ceil:
 ; CHECK:       # %bb.0:

``````````

</details>


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