nvptx "ptx-kernel" ABI (feature: abi_ptx) uses PassMode::Direct for Aggregates

[kjetilkjeka]

It produces seemingly correct behavior. However, it needs to be switched over to a less fragile ABI handling (rust-lang/compiler-team#672) before stabilization may happen.

@RalfJung's comment explains the problem and why it is considered harmless enough to allow this exception until a better solution is found.

[kjetilkjeka]

@rustbot label +O-NVPTX

[RalfJung]

@RalfJung's #117247 (comment) explains the problem and why it is considered harmless enough to allow this exception until a better solution is found.

This comment is an important part of that argument as well.

[bjorn3]

Would it be possible to use PassMode::Cast? Maybe by adapting the code at

rust/compiler/rustc_codegen_llvm/src/type_of.rs

Lines 15 to 140 in 3d1dba8

	fn uncached_llvm_type<'a, 'tcx>(
	cx: &CodegenCx<'a, 'tcx>,
	layout: TyAndLayout<'tcx>,
	defer: &mut Option<(&'a Type, TyAndLayout<'tcx>)>,
	) -> &'a Type {
	match layout.backend_repr {
	BackendRepr::Scalar(_) => bug!("handled elsewhere"),
	BackendRepr::Vector { element, count } => {
	let element = layout.scalar_llvm_type_at(cx, element);
	return cx.type_vector(element, count);
	}
	BackendRepr::Uninhabited | BackendRepr::Memory { .. } | BackendRepr::ScalarPair(..) => {}
	}
	let name = match layout.ty.kind() {
	// FIXME(eddyb) producing readable type names for trait objects can result
	// in problematically distinct types due to HRTB and subtyping (see #47638).
	// ty::Dynamic(..) |
	ty::Adt(..) | ty::Closure(..) | ty::CoroutineClosure(..) | ty::Foreign(..) | ty::Coroutine(..) | ty::Str
	// For performance reasons we use names only when emitting LLVM IR.
	if !cx.sess().fewer_names() =>
	{
	let mut name = with_no_visible_paths!(with_no_trimmed_paths!(layout.ty.to_string()));
	if let (&ty::Adt(def, _), &Variants::Single { index }) =
	(layout.ty.kind(), &layout.variants)
	{
	if def.is_enum() && !def.variants().is_empty() {
	write!(&mut name, "::{}", def.variant(index).name).unwrap();
	}
	}
	if let (&ty::Coroutine(_, _), &Variants::Single { index }) =
	(layout.ty.kind(), &layout.variants)
	{
	write!(&mut name, "::{}", ty::CoroutineArgs::variant_name(index)).unwrap();
	}
	Some(name)
	}
	_ => None,
	};
	match layout.fields {
	FieldsShape::Primitive | FieldsShape::Union(_) => {
	let fill = cx.type_padding_filler(layout.size, layout.align.abi);
	let packed = false;
	match name {
	None => cx.type_struct(&[fill], packed),
	Some(ref name) => {
	let llty = cx.type_named_struct(name);
	cx.set_struct_body(llty, &[fill], packed);
	llty
	}
	}
	}
	FieldsShape::Array { count, .. } => cx.type_array(layout.field(cx, 0).llvm_type(cx), count),
	FieldsShape::Arbitrary { .. } => match name {
	None => {
	let (llfields, packed) = struct_llfields(cx, layout);
	cx.type_struct(&llfields, packed)
	}
	Some(ref name) => {
	let llty = cx.type_named_struct(name);
	*defer = Some((llty, layout));
	llty
	}
	},
	}
	}
	fn struct_llfields<'a, 'tcx>(
	cx: &CodegenCx<'a, 'tcx>,
	layout: TyAndLayout<'tcx>,
	) -> (Vec<&'a Type>, bool) {
	debug!("struct_llfields: {:#?}", layout);
	let field_count = layout.fields.count();
	let mut packed = false;
	let mut offset = Size::ZERO;
	let mut prev_effective_align = layout.align.abi;
	let mut result: Vec<_> = Vec::with_capacity(1 + field_count * 2);
	for i in layout.fields.index_by_increasing_offset() {
	let target_offset = layout.fields.offset(i as usize);
	let field = layout.field(cx, i);
	let effective_field_align =
	layout.align.abi.min(field.align.abi).restrict_for_offset(target_offset);
	packed |= effective_field_align < field.align.abi;
	debug!(
	"struct_llfields: {}: {:?} offset: {:?} target_offset: {:?} \
	effective_field_align: {}",
	i,
	field,
	offset,
	target_offset,
	effective_field_align.bytes()
	);
	assert!(target_offset >= offset);
	let padding = target_offset - offset;
	if padding != Size::ZERO {
	let padding_align = prev_effective_align.min(effective_field_align);
	assert_eq!(offset.align_to(padding_align) + padding, target_offset);
	result.push(cx.type_padding_filler(padding, padding_align));
	debug!(" padding before: {:?}", padding);
	}
	result.push(field.llvm_type(cx));
	offset = target_offset + field.size;
	prev_effective_align = effective_field_align;
	}
	if layout.is_sized() && field_count > 0 {
	if offset > layout.size {
	bug!("layout: {:#?} stride: {:?} offset: {:?}", layout, layout.size, offset);
	}
	let padding = layout.size - offset;
	if padding != Size::ZERO {
	let padding_align = prev_effective_align;
	assert_eq!(offset.align_to(padding_align) + padding, layout.size);
	debug!(
	"struct_llfields: pad_bytes: {:?} offset: {:?} stride: {:?}",
	padding, offset, layout.size
	);
	result.push(cx.type_padding_filler(padding, padding_align));
	}
	} else {
	debug!("struct_llfields: offset: {:?} stride: {:?}", offset, layout.size);
	}
	(result, packed)
	}

(which is what effectively defines how PassMode::Direct arguments are passed currently)

[bjorn3]

@kjetilkjeka @RDambrosio016 do you think the above is feasible? Is there any documentation about what kinds of types extern "ptx-kernel" is supposed to accept? And would it be possible to eg deny anything but primitive types (including vectors) and structs?

[kjetilkjeka]

Sorry! The first message slipped my mind before I had a chance to sit down and understand it.

do you think the above is feasible?

I have heard that it's possible to achieve the same thing with PassMode::Cast. I will need to take a look at the code and experiment. But I guess you're already on this in #133932 ? Maybe I'm too late to be helpful now, but let me know if I can assist with anything other than looking at the PR.

Is there any documentation about what kinds of types extern "ptx-kernel" is supposed to accept?

I think the interop guidelines is a good place to start.

And would it be possible to eg deny anything but primitive types (including vectors) and structs?

Primitive types, vectors and structs should cover it for now.

I suspect there's some "problematic" cases of depending on the slice repr. I originally made it pass slices as two seperate argument in an attempt of being compatible with rustc_codegen_nvvm. I agree that it needs to be prohibited if we're ever going to have a chance of stabilizing a minimal version of ptx-kernel.

In the future I believe we want to provide some support for using enums, tuples, etc when calling Rust kernels from Rust host code using the same compiler. That will require design work and RFCs and is neither possible (due to building on more basic features) or a priority right now.

[bjorn3]

But I guess you're already on this in #133932 ?

Yes

Maybe I'm too late to be helpful now, but let me know if I can assist with anything other than looking at the PR.

Verifying that my PR is correct is enough for now.