[llvm] [NVPTX] Implement `isTruncateFree(EVT FromVT, EVT ToVT)` (PR #138605)

[Justin Fargnoli via llvm-commits]

https://github.com/justinfargnoli created https://github.com/llvm/llvm-project/pull/138605

This PR also makes NFC changes to `isTruncateFree(Type *SrcTy, Type *DstTy)` so that it models the HW more accurately. 

>From 5e1852a62e5b1172d2a85b6a9aec9f14b7fec15b Mon Sep 17 00:00:00 2001
From: Justin Fargnoli <jfargnoli at nvidia.com>
Date: Mon, 5 May 2025 22:49:37 +0000
Subject: [PATCH] Implement isTruncateFree(EVT FromVT, EVT ToVT)

---
 llvm/lib/Target/NVPTX/NVPTXISelLowering.h | 16 ++++++++++++----
 llvm/test/CodeGen/NVPTX/i128-array.ll     | 12 ++++++------
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.h b/llvm/lib/Target/NVPTX/NVPTXISelLowering.h
index 7a8bf3bf33a94..680ff13d8f936 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.h
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.h
@@ -155,11 +155,19 @@ class NVPTXTargetLowering : public TargetLowering {
                              Instruction *I = nullptr) const override;
 
   bool isTruncateFree(Type *SrcTy, Type *DstTy) const override {
-    // Truncating 64-bit to 32-bit is free in SASS.
-    if (!SrcTy->isIntegerTy() || !DstTy->isIntegerTy())
+    if (!(SrcTy->isIntegerTy() && DstTy->isIntegerTy()))
       return false;
-    return SrcTy->getPrimitiveSizeInBits() == 64 &&
-           DstTy->getPrimitiveSizeInBits() == 32;
+    if (SrcTy->getPrimitiveSizeInBits() <= DstTy->getPrimitiveSizeInBits())
+      return false;
+    return DstTy->getPrimitiveSizeInBits() % 32 == 0;
+  }
+
+  bool isTruncateFree(EVT FromVT, EVT ToVT) const override {
+    if (!(FromVT.isScalarInteger() && ToVT.isScalarInteger()))
+      return false;
+    if (FromVT.getSizeInBits() <= ToVT.getSizeInBits())
+      return false;
+    return ToVT.getSizeInBits() % 32 == 0;
   }
 
   EVT getSetCCResultType(const DataLayout &DL, LLVMContext &Ctx,
diff --git a/llvm/test/CodeGen/NVPTX/i128-array.ll b/llvm/test/CodeGen/NVPTX/i128-array.ll
index dd6d48bd5862c..f25d451590bed 100644
--- a/llvm/test/CodeGen/NVPTX/i128-array.ll
+++ b/llvm/test/CodeGen/NVPTX/i128-array.ll
@@ -8,13 +8,13 @@ define [2 x i128] @foo(i64 %a, i32 %b) {
 ; CHECK-NEXT:    .reg .b64 %rd<5>;
 ; CHECK-EMPTY:
 ; CHECK-NEXT:  // %bb.0:
-; CHECK-NEXT:    ld.param.u32 %r1, [foo_param_1];
 ; CHECK-NEXT:    ld.param.u64 %rd1, [foo_param_0];
-; CHECK-NEXT:    shr.s64 %rd2, %rd1, 63;
-; CHECK-NEXT:    cvt.s64.s32 %rd3, %r1;
-; CHECK-NEXT:    shr.s64 %rd4, %rd3, 63;
-; CHECK-NEXT:    st.param.v2.b64 [func_retval0], {%rd1, %rd2};
-; CHECK-NEXT:    st.param.v2.b64 [func_retval0+16], {%rd3, %rd4};
+; CHECK-NEXT:    ld.param.s32 %rd2, [foo_param_1];
+; CHECK-NEXT:    cvt.u32.u64 %r1, %rd2;
+; CHECK-NEXT:    shr.s64 %rd3, %rd1, 63;
+; CHECK-NEXT:    shr.s64 %rd4, %rd2, 63;
+; CHECK-NEXT:    st.param.v2.b64 [func_retval0], {%rd1, %rd3};
+; CHECK-NEXT:    st.param.v2.b64 [func_retval0+16], {%rd2, %rd4};
 ; CHECK-NEXT:    ret;
   %1 = sext i64 %a to i128
   %2 = sext i32 %b to i128