From 5a051bcb5e7b3cf3cf8772d057c585ca5172b7d0 Mon Sep 17 00:00:00 2001
From: tangaac <tangyan01@loongson.cn>
Date: Thu, 11 Sep 2025 10:03:36 +0800
Subject: [PATCH] Custom implement getShuffleCost

---
 .../LoongArch/LoongArchTargetTransformInfo.cpp | 18 ++++++++++++++++++
 .../LoongArch/LoongArchTargetTransformInfo.h   |  5 +++++
 2 files changed, 23 insertions(+)

diff --git a/llvm/lib/Target/LoongArch/LoongArchTargetTransformInfo.cpp b/llvm/lib/Target/LoongArch/LoongArchTargetTransformInfo.cpp
index f548a8dd0532b..17cacd49349c1 100644
--- a/llvm/lib/Target/LoongArch/LoongArchTargetTransformInfo.cpp
+++ b/llvm/lib/Target/LoongArch/LoongArchTargetTransformInfo.cpp
@@ -111,4 +111,22 @@ bool LoongArchTTIImpl::shouldExpandReduction(const IntrinsicInst *II) const {
   }
 }
 
+InstructionCost LoongArchTTIImpl::getShuffleCost(
+    TTI::ShuffleKind Kind, VectorType *DstTy, VectorType *SrcTy,
+    ArrayRef<int> Mask, TTI::TargetCostKind CostKind, int Index,
+    VectorType *SubTp, ArrayRef<const Value *> Args,
+    const Instruction *CxtI) const {
+
+  Kind = improveShuffleKindFromMask(Kind, Mask, SrcTy, Index, SubTp);
+  switch (Kind) {
+  default:
+    return BaseT::getShuffleCost(Kind, DstTy, SrcTy, Mask, CostKind, Index,
+                                 SubTp, Args, CxtI);
+  // case TTI::SK_Broadcast:
+  // case TTI::SK_Reverse:
+  case TTI::SK_Transpose:
+    return 1;
+  }
+}
+
 // TODO: Implement more hooks to provide TTI machinery for LoongArch.
diff --git a/llvm/lib/Target/LoongArch/LoongArchTargetTransformInfo.h b/llvm/lib/Target/LoongArch/LoongArchTargetTransformInfo.h
index e3f16c7804994..5a36c34735174 100644
--- a/llvm/lib/Target/LoongArch/LoongArchTargetTransformInfo.h
+++ b/llvm/lib/Target/LoongArch/LoongArchTargetTransformInfo.h
@@ -54,6 +54,11 @@ class LoongArchTTIImpl : public BasicTTIImplBase<LoongArchTTIImpl> {
   bool enableWritePrefetching() const override;
 
   bool shouldExpandReduction(const IntrinsicInst *II) const override;
+  InstructionCost
+  getShuffleCost(TTI::ShuffleKind Kind, VectorType *DstTy, VectorType *SrcTy,
+                 ArrayRef<int> Mask, TTI::TargetCostKind CostKind, int Index,
+                 VectorType *SubTp, ArrayRef<const Value *> Args = {},
+                 const Instruction *CxtI = nullptr) const override;
 
   // TODO: Implement more hooks to provide TTI machinery for LoongArch.
 };