From 14348d951912022c4b23d7316f06a78a2791da01 Mon Sep 17 00:00:00 2001
From: Kjetil Kjeka <kjetil@muybridge.com>
Date: Fri, 31 May 2024 22:45:27 +0200
Subject: [PATCH] NVPTX: Avoid PassMode::Direct for C ABI

---
 compiler/rustc_target/src/abi/call/nvptx64.rs | 47 +++++++++++++++----
 tests/assembly/nvptx-c-abi-ret-v7.rs          |  1 -
 tests/ui/abi/compatibility.rs                 | 10 ++--
 3 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/compiler/rustc_target/src/abi/call/nvptx64.rs b/compiler/rustc_target/src/abi/call/nvptx64.rs
index c6b6f66e6be..ac68e8879f6 100644
--- a/compiler/rustc_target/src/abi/call/nvptx64.rs
+++ b/compiler/rustc_target/src/abi/call/nvptx64.rs
@@ -1,23 +1,54 @@
 use crate::abi::call::{ArgAbi, FnAbi, PassMode, Reg, Size, Uniform};
 use crate::abi::{HasDataLayout, TyAbiInterface};
 
+use super::{ArgAttribute, ArgAttributes, ArgExtension, CastTarget};
+
 fn classify_ret<Ty>(ret: &mut ArgAbi<'_, Ty>) {
-    if ret.layout.is_aggregate() && ret.layout.size.bits() > 64 {
-        ret.make_indirect();
-    } else {
-        // FIXME: this is wrong! Need to decide which ABI we really want here.
-        ret.make_direct_deprecated();
+    if ret.layout.is_aggregate() && ret.layout.is_sized() {
+        classify_aggregate(ret)
+    } else if ret.layout.size.bits() < 32 && ret.layout.is_sized() {
+        ret.extend_integer_width_to(32);
     }
 }
 
 fn classify_arg<Ty>(arg: &mut ArgAbi<'_, Ty>) {
-    if arg.layout.is_aggregate() {
-        arg.make_indirect_byval(None);
-    } else if arg.layout.size.bits() < 32 {
+    if arg.layout.is_aggregate() && arg.layout.is_sized() {
+        classify_aggregate(arg)
+    } else if arg.layout.size.bits() < 32 && arg.layout.is_sized() {
         arg.extend_integer_width_to(32);
     }
 }
 
+/// the pass mode used for aggregates in arg and ret position
+fn classify_aggregate<Ty>(arg: &mut ArgAbi<'_, Ty>) {
+    let align_bytes = arg.layout.align.abi.bytes();
+    let size = arg.layout.size;
+
+    let reg = match align_bytes {
+        1 => Reg::i8(),
+        2 => Reg::i16(),
+        4 => Reg::i32(),
+        8 => Reg::i64(),
+        16 => Reg::i128(),
+        _ => unreachable!("Align is given as power of 2 no larger than 16 bytes"),
+    };
+
+    if align_bytes == size.bytes() {
+        arg.cast_to(CastTarget {
+            prefix: [Some(reg), None, None, None, None, None, None, None],
+            rest: Uniform::new(Reg::i8(), Size::from_bytes(0)),
+            attrs: ArgAttributes {
+                regular: ArgAttribute::default(),
+                arg_ext: ArgExtension::None,
+                pointee_size: Size::ZERO,
+                pointee_align: None,
+            },
+        });
+    } else {
+        arg.cast_to(Uniform::new(reg, size));
+    }
+}
+
 fn classify_arg_kernel<'a, Ty, C>(_cx: &C, arg: &mut ArgAbi<'a, Ty>)
 where
     Ty: TyAbiInterface<'a, C> + Copy,
diff --git a/tests/assembly/nvptx-c-abi-ret-v7.rs b/tests/assembly/nvptx-c-abi-ret-v7.rs
index 0c1594a5f20..8992cb8d540 100644
--- a/tests/assembly/nvptx-c-abi-ret-v7.rs
+++ b/tests/assembly/nvptx-c-abi-ret-v7.rs
@@ -1,7 +1,6 @@
 //@ assembly-output: ptx-linker
 //@ compile-flags: --crate-type cdylib -C target-cpu=sm_86 -Z unstable-options -Clinker-flavor=llbc
 //@ only-nvptx64
-//@ ignore-nvptx64
 
 // The PTX ABI stability is tied to major versions of the PTX ISA
 // These tests assume major version 7
diff --git a/tests/ui/abi/compatibility.rs b/tests/ui/abi/compatibility.rs
index c27a107a639..1be9ada2822 100644
--- a/tests/ui/abi/compatibility.rs
+++ b/tests/ui/abi/compatibility.rs
@@ -55,13 +55,9 @@
 //@ revisions: csky
 //@[csky] compile-flags: --target csky-unknown-linux-gnuabiv2
 //@[csky] needs-llvm-components: csky
-
-// FIXME: disabled on nvptx64 since the target ABI fails the sanity check
-// see https://github.com/rust-lang/rust/issues/117480
-/* revisions: nvptx64
-  [nvptx64] compile-flags: --target nvptx64-nvidia-cuda
-  [nvptx64] needs-llvm-components: nvptx
-*/
+//@ revisions: nvptx64
+//@[nvptx64] compile-flags: --target nvptx64-nvidia-cuda
+//@[nvptx64] needs-llvm-components: nvptx
 #![feature(rustc_attrs, unsized_fn_params, transparent_unions)]
 #![cfg_attr(not(host), feature(no_core, lang_items), no_std, no_core)]
 #![allow(unused, improper_ctypes_definitions, internal_features)]