Emit alignment for all memory ops

Author: jwollen

crates/rustc_codegen_spirv/src/builder/builder_methods.rs

@@ -1473,11 +1473,17 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
         self.fatal("dynamic alloca not supported yet")
     }
 
-    fn load(&mut self, ty: Self::Type, ptr: Self::Value, _align: Align) -> Self::Value {
+    fn load(&mut self, ty: Self::Type, ptr: Self::Value, align: Align) -> Self::Value {
         let (ptr, access_ty) = self.adjust_pointer_for_typed_access(ptr, ty);
         let loaded_val = ptr.const_fold_load(self).unwrap_or_else(|| {
             self.emit()
-                .load(access_ty, None, ptr.def(self), None, empty())
+                .load(
+                    access_ty,
+                    None,
+                    ptr.def(self),
+                    Some(MemoryAccess::ALIGNED),
+                    std::iter::once(Operand::LiteralBit32(align.bytes() as _)),
+                )
                 .unwrap()
                 .with_type(access_ty)
         });
@@ -1599,12 +1605,17 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
         // ignore
     }
 
-    fn store(&mut self, val: Self::Value, ptr: Self::Value, _align: Align) -> Self::Value {
+    fn store(&mut self, val: Self::Value, ptr: Self::Value, align: Align) -> Self::Value {
         let (ptr, access_ty) = self.adjust_pointer_for_typed_access(ptr, val.ty);
         let val = self.bitcast(val, access_ty);
 
         self.emit()
-            .store(ptr.def(self), val.def(self), None, empty())
+            .store(
+                ptr.def(self),
+                val.def(self),
+                Some(MemoryAccess::ALIGNED),
+                std::iter::once(Operand::LiteralBit32(align.bytes() as _)),
+            )
             .unwrap();
         // FIXME(eddyb) this is meant to be a handle the store instruction itself.
         val
@@ -2262,9 +2273,9 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
     fn memcpy(
         &mut self,
         dst: Self::Value,
-        _dst_align: Align,
+        dst_align: Align,
         src: Self::Value,
-        _src_align: Align,
+        src_align: Align,
         size: Self::Value,
         flags: MemFlags,
     ) {
@@ -2302,12 +2313,29 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
             }
         });
 
+        // Pass all operands as `additional_params` since rspirv doesn't allow specifying
+        // extra operands ofter the first `MemoryAccess`
+        let mut ops: SmallVec<[_; 4]> = Default::default();
+        ops.push(Operand::MemoryAccess(MemoryAccess::ALIGNED));
+        if src_align != dst_align {
+            if self.emit().version().unwrap() > (1, 3) {
+                ops.push(Operand::LiteralBit32(dst_align.bytes() as _));
+                ops.push(Operand::MemoryAccess(MemoryAccess::ALIGNED));
+                ops.push(Operand::LiteralBit32(src_align.bytes() as _));
+            } else {
+                let align = dst_align.min(src_align);
+                ops.push(Operand::LiteralBit32(align.bytes() as _));
+            }
+        } else {
+            ops.push(Operand::LiteralBit32(dst_align.bytes() as _));
+        }
+
         if let Some((dst, src)) = typed_copy_dst_src {
             if let Some(const_value) = src.const_fold_load(self) {
                 self.store(const_value, dst, Align::from_bytes(0).unwrap());
             } else {
                 self.emit()
-                    .copy_memory(dst.def(self), src.def(self), None, None, empty())
+                    .copy_memory(dst.def(self), src.def(self), None, None, ops)
                     .unwrap();
             }
         } else {
@@ -2318,7 +2346,7 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
                     size.def(self),
                     None,
                     None,
-                    empty(),
+                    ops,
                 )
                 .unwrap();
             self.zombie(dst.def(self), "cannot memcpy dynamically sized data");

crates/rustc_codegen_spirv/src/linker/mem2reg.rs

@@ -11,11 +11,12 @@
 
 use super::simple_passes::outgoing_edges;
 use super::{apply_rewrite_rules, id};
+use itertools::Itertools;
 use rspirv::dr::{Block, Function, Instruction, ModuleHeader, Operand};
 use rspirv::spirv::{Op, Word};
 use rustc_data_structures::fx::{FxHashMap, FxHashSet, FxIndexMap};
-use rustc_middle::bug;
 use std::collections::hash_map;
+use std::iter;
 
 // HACK(eddyb) newtype instead of type alias to avoid mistakes.
 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
@@ -328,10 +329,15 @@ fn split_copy_memory(
             if inst.class.opcode == Op::CopyMemory {
                 let target = inst.operands[0].id_ref_any().unwrap();
                 let source = inst.operands[1].id_ref_any().unwrap();
-                if inst.operands.len() > 2 {
-                    // TODO: Copy the memory operands to the load/store
-                    bug!("mem2reg OpCopyMemory doesn't support memory operands yet");
-                }
+                let mem_ops = &inst.operands[2..];
+                let (store_mem_ops, load_mem_ops) = if let Some((index, _)) = mem_ops[1..]
+                    .iter()
+                    .find_position(|op| matches!(op, Operand::MemoryAccess(..)))
+                {
+                    mem_ops.split_at(index)
+                } else {
+                    (mem_ops, mem_ops)
+                };
                 let ty = match (var_map.get(&target), var_map.get(&source)) {
                     (None, None) => {
                         inst_index += 1;
@@ -345,17 +351,22 @@ fn split_copy_memory(
                     }
                 };
                 let temp_id = id(header);
+
+                let load_ops = iter::once(Operand::IdRef(source))
+                    .chain(load_mem_ops.iter().cloned())
+                    .collect();
+
+                let store_ops = [Operand::IdRef(target), Operand::IdRef(temp_id)]
+                    .into_iter()
+                    .chain(store_mem_ops.iter().cloned())
+                    .collect();
+
                 block.instructions[inst_index] =
-                    Instruction::new(Op::Load, Some(ty), Some(temp_id), vec![Operand::IdRef(
-                        source,
-                    )]);
+                    Instruction::new(Op::Load, Some(ty), Some(temp_id), load_ops);
                 inst_index += 1;
                 block.instructions.insert(
                     inst_index,
-                    Instruction::new(Op::Store, None, None, vec![
-                        Operand::IdRef(target),
-                        Operand::IdRef(temp_id),
-                    ]),
+                    Instruction::new(Op::Store, None, None, store_ops),
                 );
             }
             inst_index += 1;