[llvm] [NVPTX] Allow more argument integer types, such as i256 and i96 (PR #154824)

[Alex MacLean via llvm-commits]

https://github.com/AlexMaclean updated https://github.com/llvm/llvm-project/pull/154824

>From f56c87f19acbefbe91c99bc248bcc4f4b86bc4a7 Mon Sep 17 00:00:00 2001
From: Alex Maclean <amaclean at nvidia.com>
Date: Thu, 21 Aug 2025 19:05:39 +0000
Subject: [PATCH] [NVPTX] Allow more argument integer types, such as i256 and
 i96

---
 llvm/lib/Target/NVPTX/NVPTXUtilities.h             | 4 +++-
 llvm/test/CodeGen/Generic/add-with-overflow-128.ll | 3 ---
 llvm/test/CodeGen/Generic/icmp-illegal.ll          | 4 ----
 3 files changed, 3 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Target/NVPTX/NVPTXUtilities.h b/llvm/lib/Target/NVPTX/NVPTXUtilities.h
index 9421f9f54d0a6..a070789f85e0b 100644
--- a/llvm/lib/Target/NVPTX/NVPTXUtilities.h
+++ b/llvm/lib/Target/NVPTX/NVPTXUtilities.h
@@ -82,6 +82,8 @@ inline unsigned promoteScalarArgumentSize(unsigned size) {
     return 32;
   if (size <= 64)
     return 64;
+  if (size <= 128)
+    return 128;
   return size;
 }
 
@@ -89,7 +91,7 @@ bool shouldEmitPTXNoReturn(const Value *V, const TargetMachine &TM);
 
 inline bool shouldPassAsArray(Type *Ty) {
   return Ty->isAggregateType() || Ty->isVectorTy() ||
-         Ty->getScalarSizeInBits() == 128 || Ty->isHalfTy() || Ty->isBFloatTy();
+         Ty->getScalarSizeInBits() >= 128 || Ty->isHalfTy() || Ty->isBFloatTy();
 }
 
 namespace NVPTX {
diff --git a/llvm/test/CodeGen/Generic/add-with-overflow-128.ll b/llvm/test/CodeGen/Generic/add-with-overflow-128.ll
index 389e6511b6c7c..3db446663e71d 100644
--- a/llvm/test/CodeGen/Generic/add-with-overflow-128.ll
+++ b/llvm/test/CodeGen/Generic/add-with-overflow-128.ll
@@ -1,8 +1,5 @@
 ; RUN: llc < %s
 
-; NVPTX fails to LowerFormalArguments for arg type i96
-; the arg byte size must be one of the {16, 8, 4, 2}
-; XFAIL: target=nvptx{{.*}}
 
 @ok = internal constant [4 x i8] c"%d\0A\00"
 @no = internal constant [4 x i8] c"no\0A\00"
diff --git a/llvm/test/CodeGen/Generic/icmp-illegal.ll b/llvm/test/CodeGen/Generic/icmp-illegal.ll
index 1af42e3443d21..77dd5a59dfd1e 100644
--- a/llvm/test/CodeGen/Generic/icmp-illegal.ll
+++ b/llvm/test/CodeGen/Generic/icmp-illegal.ll
@@ -1,9 +1,5 @@
 ; RUN: llc < %s | FileCheck %s
 
-; NVPTX fails to LowerFormalArguments for arg size > i128
-; the arg byte size must be one of the {16, 8, 4, 2}
-; XFAIL: target=nvptx{{.*}}
-
 ; CHECK-LABEL: test_ult
 define i1 @test_ult(i256 %a) nounwind {
   %1 = icmp ult i256 %a, -6432394258550908438