[llvm] [LoongArch] Lower ANY_EXTEND to ZERO_EXTEND under LASX to avoid scalarization (PR #201099)

[via llvm-commits]

https://github.com/lrzlin updated https://github.com/llvm/llvm-project/pull/201099

>From c38f00418edd43dc4f822fdb900bc51a592916b0 Mon Sep 17 00:00:00 2001
From: Lin Runze <linrunze at loongson.cn>
Date: Tue, 26 May 2026 17:59:30 +0800
Subject: [PATCH] [LoongArch] Lower ANY_EXTEND to ZERO_EXTEND under LASX to
 avoid scalarization

---
 llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp | 12 ++++++++++++
 llvm/lib/Target/LoongArch/LoongArchISelLowering.h   |  1 +
 2 files changed, 13 insertions(+)

diff --git a/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp b/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp
index 561064b3e1090..79a4ce24c09b0 100644
--- a/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp
+++ b/llvm/lib/Target/LoongArch/LoongArchISelLowering.cpp
@@ -487,6 +487,7 @@ LoongArchTargetLowering::LoongArchTargetLowering(const TargetMachine &TM,
     for (MVT VT : {MVT::v4i64, MVT::v8i32, MVT::v16i16}) {
       setOperationAction(ISD::SIGN_EXTEND, VT, Legal);
       setOperationAction(ISD::ZERO_EXTEND, VT, Legal);
+      setOperationAction(ISD::ANY_EXTEND, VT, Custom);
     }
     for (MVT VT :
          {MVT::v2i64, MVT::v4i32, MVT::v4i64, MVT::v8i16, MVT::v8i32}) {
@@ -661,6 +662,8 @@ SDValue LoongArchTargetLowering::LowerOperation(SDValue Op,
     return lowerSIGN_EXTEND_VECTOR_INREG(Op, DAG);
   case ISD::DYNAMIC_STACKALLOC:
     return lowerDYNAMIC_STACKALLOC(Op, DAG);
+  case ISD::ANY_EXTEND:
+    return lowerANY_EXTEND(Op, DAG);
   }
   return SDValue();
 }
@@ -1010,6 +1013,15 @@ SDValue LoongArchTargetLowering::lowerSIGN_EXTEND_VECTOR_INREG(
   return SDValue();
 }
 
+// ANY_EXTEND can be replaced by ZERO_EXTEND when LASX is enabled.
+SDValue LoongArchTargetLowering::lowerANY_EXTEND(SDValue Op,
+                                                 SelectionDAG &DAG) const {
+  SDLoc DL(Op);
+  if (!Subtarget.hasExtLASX() || !Subtarget.is64Bit())
+    return SDValue();
+  return DAG.getNode(ISD::ZERO_EXTEND, DL, Op.getValueType(), Op.getOperand(0));
+}
+
 // Lower vecreduce_add using vhaddw instructions.
 // For Example:
 //  call i32 @llvm.vector.reduce.add.v4i32(<4 x i32> %a)
diff --git a/llvm/lib/Target/LoongArch/LoongArchISelLowering.h b/llvm/lib/Target/LoongArch/LoongArchISelLowering.h
index 189ecbe4820d2..47589706f62fc 100644
--- a/llvm/lib/Target/LoongArch/LoongArchISelLowering.h
+++ b/llvm/lib/Target/LoongArch/LoongArchISelLowering.h
@@ -259,6 +259,7 @@ class LoongArchTargetLowering : public TargetLowering {
   SDValue lowerFP_EXTEND(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerSIGN_EXTEND_VECTOR_INREG(SDValue Op, SelectionDAG &DAG) const;
   SDValue lowerDYNAMIC_STACKALLOC(SDValue Op, SelectionDAG &DAG) const;
+  SDValue lowerANY_EXTEND(SDValue Op, SelectionDAG &DAG) const;
 
   bool isFPImmLegal(const APFloat &Imm, EVT VT,
                     bool ForCodeSize) const override;