Access the original type's fields through a pointercast, under the Logical addressing model. (#469)

* Defer pointer casts under the Logical addressing model.

* Access the original type's fields through a `pointercast`, under the Logical addressing model.

* Add a test using a `for` loop and a custom `Range`-like iterator.

Author: eddyb

crates/rustc_codegen_spirv/src/builder/builder_methods.rs

@@ -1,10 +1,8 @@
 use super::Builder;
-use crate::builder_spirv::{BuilderCursor, SpirvValue, SpirvValueExt};
+use crate::builder_spirv::{BuilderCursor, SpirvValue, SpirvValueExt, SpirvValueKind};
 use crate::spirv_type::SpirvType;
 use rspirv::dr::{InsertPoint, Instruction, Operand};
-use rspirv::spirv::{
-    AddressingModel, Capability, MemoryModel, MemorySemantics, Op, Scope, StorageClass, Word,
-};
+use rspirv::spirv::{Capability, MemoryModel, MemorySemantics, Op, Scope, StorageClass, Word};
 use rustc_codegen_ssa::common::{
     AtomicOrdering, AtomicRmwBinOp, IntPredicate, RealPredicate, SynchronizationScope,
 };
@@ -18,6 +16,7 @@ use rustc_middle::bug;
 use rustc_middle::ty::Ty;
 use rustc_span::Span;
 use rustc_target::abi::{Abi, Align, Scalar, Size};
+use std::convert::TryInto;
 use std::iter::empty;
 use std::ops::Range;
 
@@ -334,54 +333,74 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         }
     }
 
-    fn zombie_bitcast_ptr(&self, def: Word, from_ty: Word, to_ty: Word) {
-        let is_logical = self
-            .emit()
-            .module_ref()
-            .memory_model
-            .as_ref()
-            .map_or(false, |inst| {
-                inst.operands[0].unwrap_addressing_model() == AddressingModel::Logical
-            });
-        if is_logical {
-            if self.is_system_crate() {
-                self.zombie(def, "OpBitcast on ptr without AddressingModel != Logical")
-            } else {
-                self.struct_err("Cannot cast between pointer types")
-                    .note(&format!("from: {}", self.debug_type(from_ty)))
-                    .note(&format!("to: {}", self.debug_type(to_ty)))
-                    .emit()
-            }
-        }
-    }
+    /// If possible, return the appropriate `OpAccessChain` indices for going from
+    /// a pointer to `ty`, to a pointer to `leaf_ty`, with an added `offset`.
+    ///
+    /// That is, try to turn `((_: *T) as *u8).add(offset) as *Leaf` into a series
+    /// of struct field and array/vector element accesses.
+    fn recover_access_chain_from_offset(
+        &self,
+        mut ty: Word,
+        leaf_ty: Word,
+        mut offset: Size,
+    ) -> Option<Vec<u32>> {
+        assert_ne!(ty, leaf_ty);
+
+        // NOTE(eddyb) `ty` and `ty_kind` should be kept in sync.
+        let mut ty_kind = self.lookup_type(ty);
 
-    // Sometimes, when accessing the first field of a struct, vector, etc., instead of calling
-    // struct_gep, codegen_ssa will call pointercast. This will then try to catch those cases and
-    // translate them back to a struct_gep, instead of failing to compile the OpBitcast (which is
-    // unsupported on shader target)
-    fn try_pointercast_via_gep(&self, mut val: Word, field: Word) -> Option<Vec<u32>> {
         let mut indices = Vec::new();
-        while val != field {
-            match self.lookup_type(val) {
+        loop {
+            match ty_kind {
                 SpirvType::Adt {
                     field_types,
                     field_offsets,
                     ..
                 } => {
-                    let index = field_offsets.iter().position(|&off| off == Size::ZERO)?;
-                    indices.push(index as u32);
-                    val = field_types[index];
+                    let (i, field_ty, field_ty_kind, offset_in_field) = field_offsets
+                        .iter()
+                        .enumerate()
+                        .find_map(|(i, &field_offset)| {
+                            if field_offset > offset {
+                                return None;
+                            }
+
+                            // Grab the actual field type to be able to confirm that
+                            // the leaf is somewhere inside the field.
+                            let field_ty = field_types[i];
+                            let field_ty_kind = self.lookup_type(field_ty);
+
+                            let offset_in_field = offset - field_offset;
+                            if offset_in_field < field_ty_kind.sizeof(self)? {
+                                Some((i, field_ty, field_ty_kind, offset_in_field))
+                            } else {
+                                None
+                            }
+                        })?;
+
+                    ty = field_ty;
+                    ty_kind = field_ty_kind;
+
+                    indices.push(i as u32);
+                    offset = offset_in_field;
                 }
                 SpirvType::Vector { element, .. }
                 | SpirvType::Array { element, .. }
                 | SpirvType::RuntimeArray { element } => {
-                    indices.push(0);
-                    val = element;
+                    ty = element;
+                    ty_kind = self.lookup_type(ty);
+
+                    let stride = ty_kind.sizeof(self)?;
+                    indices.push((offset.bytes() / stride.bytes()).try_into().ok()?);
+                    offset = Size::from_bytes(offset.bytes() % stride.bytes());
                 }
                 _ => return None,
             }
+
+            if offset == Size::ZERO && ty == leaf_ty {
+                return Some(indices);
+            }
         }
-        Some(indices)
     }
 }
 
@@ -963,26 +982,65 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
     }
 
     fn struct_gep(&mut self, ptr: Self::Value, idx: u64) -> Self::Value {
-        let result_pointee_type = match self.lookup_type(ptr.ty) {
-            SpirvType::Pointer { pointee } => match self.lookup_type(pointee) {
-                SpirvType::Adt { field_types, .. } => field_types[idx as usize],
-                SpirvType::Array { element, .. }
-                | SpirvType::RuntimeArray { element, .. }
-                | SpirvType::Vector { element, .. } => element,
-                other => self.fatal(&format!(
-                    "struct_gep not on struct, array, or vector type: {:?}, index {}",
-                    other, idx
-                )),
-            },
+        let pointee = match self.lookup_type(ptr.ty) {
+            SpirvType::Pointer { pointee } => pointee,
             other => self.fatal(&format!(
                 "struct_gep not on pointer type: {:?}, index {}",
                 other, idx
             )),
         };
+        let pointee_kind = self.lookup_type(pointee);
+        let result_pointee_type = match pointee_kind {
+            SpirvType::Adt {
+                ref field_types, ..
+            } => field_types[idx as usize],
+            SpirvType::Array { element, .. }
+            | SpirvType::RuntimeArray { element, .. }
+            | SpirvType::Vector { element, .. } => element,
+            other => self.fatal(&format!(
+                "struct_gep not on struct, array, or vector type: {:?}, index {}",
+                other, idx
+            )),
+        };
         let result_type = SpirvType::Pointer {
             pointee: result_pointee_type,
         }
         .def(self.span(), self);
+
+        // Special-case field accesses through a `pointercast`, to accesss the
+        // right field in the original type, for the `Logical` addressing model.
+        if let SpirvValueKind::LogicalPtrCast {
+            original_ptr,
+            original_pointee_ty,
+            zombie_target_undef: _,
+        } = ptr.kind
+        {
+            let offset = match pointee_kind {
+                SpirvType::Adt { field_offsets, .. } => field_offsets[idx as usize],
+                SpirvType::Array { element, .. }
+                | SpirvType::RuntimeArray { element, .. }
+                | SpirvType::Vector { element, .. } => {
+                    self.lookup_type(element).sizeof(self).unwrap() * idx
+                }
+                _ => unreachable!(),
+            };
+            if let Some(indices) = self.recover_access_chain_from_offset(
+                original_pointee_ty,
+                result_pointee_type,
+                offset,
+            ) {
+                let indices = indices
+                    .into_iter()
+                    .map(|idx| self.constant_u32(self.span(), idx).def(self))
+                    .collect::<Vec<_>>();
+                return self
+                    .emit()
+                    .access_chain(result_type, None, original_ptr, indices)
+                    .unwrap()
+                    .with_type(result_type);
+            }
+        }
+
         // Important! LLVM, and therefore intel-compute-runtime, require the `getelementptr` instruction (and therefore
         // OpAccessChain) on structs to be a constant i32. Not i64! i32.
         if idx > u32::MAX as u64 {
@@ -1131,16 +1189,34 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
         if val.ty == dest_ty {
             val
         } else {
+            let val_is_ptr = matches!(self.lookup_type(val.ty), SpirvType::Pointer { .. });
+            let dest_is_ptr = matches!(self.lookup_type(dest_ty), SpirvType::Pointer { .. });
+
+            // Reuse the pointer-specific logic in `pointercast` for `*T -> *U`.
+            if val_is_ptr && dest_is_ptr {
+                return self.pointercast(val, dest_ty);
+            }
+
             let result = self
                 .emit()
                 .bitcast(dest_ty, None, val.def(self))
                 .unwrap()
                 .with_type(dest_ty);
-            let val_is_ptr = matches!(self.lookup_type(val.ty), SpirvType::Pointer { .. });
-            let dest_is_ptr = matches!(self.lookup_type(dest_ty), SpirvType::Pointer { .. });
-            if val_is_ptr || dest_is_ptr {
-                self.zombie_bitcast_ptr(result.def(self), val.ty, dest_ty);
+
+            if (val_is_ptr || dest_is_ptr) && self.logical_addressing_model() {
+                if self.is_system_crate() {
+                    self.zombie(
+                        result.def(self),
+                        "OpBitcast between ptr and non-ptr without AddressingModel != Logical",
+                    )
+                } else {
+                    self.struct_err("Cannot cast between pointer and non-pointer types")
+                        .note(&format!("from: {}", self.debug_type(val.ty)))
+                        .note(&format!("to: {}", self.debug_type(dest_ty)))
+                        .emit()
+                }
             }
+
             result
         }
     }
@@ -1204,12 +1280,29 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
     }
 
     fn pointercast(&mut self, val: Self::Value, dest_ty: Self::Type) -> Self::Value {
-        let val_pointee = match self.lookup_type(val.ty) {
-            SpirvType::Pointer { pointee } => pointee,
-            other => self.fatal(&format!(
-                "pointercast called on non-pointer source type: {:?}",
-                other
-            )),
+        let (val, val_pointee) = match val.kind {
+            // Strip a previous `pointercast`, to reveal the original pointer type.
+            SpirvValueKind::LogicalPtrCast {
+                original_ptr,
+                original_pointee_ty,
+                zombie_target_undef: _,
+            } => (
+                original_ptr.with_type(
+                    SpirvType::Pointer {
+                        pointee: original_pointee_ty,
+                    }
+                    .def(self.span(), self),
+                ),
+                original_pointee_ty,
+            ),
+
+            _ => match self.lookup_type(val.ty) {
+                SpirvType::Pointer { pointee } => (val, pointee),
+                other => self.fatal(&format!(
+                    "pointercast called on non-pointer source type: {:?}",
+                    other
+                )),
+            },
         };
         let dest_pointee = match self.lookup_type(dest_ty) {
             SpirvType::Pointer { pointee } => pointee,
@@ -1220,7 +1313,9 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
         };
         if val.ty == dest_ty {
             val
-        } else if let Some(indices) = self.try_pointercast_via_gep(val_pointee, dest_pointee) {
+        } else if let Some(indices) =
+            self.recover_access_chain_from_offset(val_pointee, dest_pointee, Size::ZERO)
+        {
             let indices = indices
                 .into_iter()
                 .map(|idx| self.constant_u32(self.span(), idx).def(self))
@@ -1229,14 +1324,21 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
                 .access_chain(dest_ty, None, val.def(self), indices)
                 .unwrap()
                 .with_type(dest_ty)
+        } else if self.logical_addressing_model() {
+            // Defer the cast so that it has a chance to be avoided.
+            SpirvValue {
+                kind: SpirvValueKind::LogicalPtrCast {
+                    original_ptr: val.def(self),
+                    original_pointee_ty: val_pointee,
+                    zombie_target_undef: self.undef(dest_ty).def(self),
+                },
+                ty: dest_ty,
+            }
         } else {
-            let result = self
-                .emit()
+            self.emit()
                 .bitcast(dest_ty, None, val.def(self))
                 .unwrap()
-                .with_type(dest_ty);
-            self.zombie_bitcast_ptr(result.def(self), val.ty, dest_ty);
-            result
+                .with_type(dest_ty)
         }
     }
 

crates/rustc_codegen_spirv/src/builder_spirv.rs

@@ -13,6 +13,7 @@ use std::{fs::File, io::Write, path::Path};
 #[derive(Copy, Clone, Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
 pub enum SpirvValueKind {
     Def(Word),
+
     /// There are a fair number of places where `rustc_codegen_ssa` creates a pointer to something
     /// that cannot be pointed to in SPIR-V. For example, constant values are frequently emitted as
     /// a pointer to constant memory, and then dereferenced where they're used. Functions are the
@@ -28,6 +29,24 @@ pub enum SpirvValueKind {
         /// its initializer) to attach zombies to.
         global_var: Word,
     },
+
+    /// Deferred pointer cast, for the `Logical` addressing model (which doesn't
+    /// really support raw pointers in the way Rust expects to be able to use).
+    ///
+    /// The cast's target pointer type is the `ty` of the `SpirvValue` that has
+    /// `LogicalPtrCast` as its `kind`, as it would be redundant to have it here.
+    LogicalPtrCast {
+        /// Pointer value being cast.
+        original_ptr: Word,
+
+        /// Pointee type of `original_ptr`.
+        original_pointee_ty: Word,
+
+        /// `OpUndef` of the right target pointer type, to attach zombies to.
+        // FIXME(eddyb) we should be using a real `OpBitcast` here, but we can't
+        // emit that on the fly during `SpirvValue::def`, due to builder locking.
+        zombie_target_undef: Word,
+    },
 }
 
 #[derive(Copy, Clone, Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
@@ -49,7 +68,8 @@ impl SpirvValue {
                 };
                 Some(initializer.with_type(ty))
             }
-            SpirvValueKind::Def(_) => None,
+
+            SpirvValueKind::Def(_) | SpirvValueKind::LogicalPtrCast { .. } => None,
         }
     }
 
@@ -69,6 +89,7 @@ impl SpirvValue {
     pub fn def_with_span(self, cx: &CodegenCx<'_>, span: Span) -> Word {
         match self.kind {
             SpirvValueKind::Def(word) => word,
+
             SpirvValueKind::ConstantPointer {
                 initializer: _,
                 global_var,
@@ -83,6 +104,29 @@ impl SpirvValue {
 
                 global_var
             }
+
+            SpirvValueKind::LogicalPtrCast {
+                original_ptr: _,
+                original_pointee_ty,
+                zombie_target_undef,
+            } => {
+                if cx.is_system_crate() {
+                    cx.zombie_with_span(
+                        zombie_target_undef,
+                        span,
+                        "OpBitcast on ptr without AddressingModel != Logical",
+                    )
+                } else {
+                    cx.tcx
+                        .sess
+                        .struct_span_err(span, "Cannot cast between pointer types")
+                        .note(&format!("from: *{}", cx.debug_type(original_pointee_ty)))
+                        .note(&format!("to: {}", cx.debug_type(self.ty)))
+                        .emit()
+                }
+
+                zombie_target_undef
+            }
         }
     }
 }