builder: constant-fold bx.mul(a, b) in order to support ptr::{read,write}. (#485)

eddyb · web-flow · commit 7348f3389d49 · 2021-03-11T09:53:49.000+01:00
diff --git a/crates/rustc_codegen_spirv/src/builder/builder_methods.rs b/crates/rustc_codegen_spirv/src/builder/builder_methods.rs
@@ -1,5 +1,5 @@
 use super::Builder;
-use crate::builder_spirv::{BuilderCursor, SpirvValue, SpirvValueExt, SpirvValueKind};
+use crate::builder_spirv::{BuilderCursor, SpirvConst, SpirvValue, SpirvValueExt, SpirvValueKind};
 use crate::spirv_type::SpirvType;
 use rspirv::dr::{InsertPoint, Instruction, Operand};
 use rspirv::spirv::{Capability, MemoryModel, MemorySemantics, Op, Scope, StorageClass, Word};
@@ -21,10 +21,44 @@ use std::iter::empty;
 use std::ops::Range;
 
 macro_rules! simple_op {
-    ($func_name:ident, $inst_name:ident) => {
+    (
+        $func_name:ident, $inst_name:ident
+        $(, fold_const {
+            $(int($fold_int_lhs:ident, $fold_int_rhs:ident) => $fold_int:expr)?
+        })?
+    ) => {
         fn $func_name(&mut self, lhs: Self::Value, rhs: Self::Value) -> Self::Value {
             assert_ty_eq!(self, lhs.ty, rhs.ty);
             let result_type = lhs.ty;
+
+            $(if let Some(const_lhs) = self.builder.lookup_const(lhs) {
+                if let Some(const_rhs) = self.builder.lookup_const(rhs) {
+                    match self.lookup_type(result_type) {
+                        $(SpirvType::Integer(bits, signed) => {
+                            let size = Size::from_bits(bits);
+                            let as_u128 = |const_val| {
+                                let x = match const_val {
+                                    SpirvConst::U32(_, x) => x as u128,
+                                    SpirvConst::U64(_, x) => x as u128,
+                                    _ => return None,
+                                };
+                                Some(if signed {
+                                    size.sign_extend(x)
+                                } else {
+                                    size.truncate(x)
+                                })
+                            };
+                            if let Some($fold_int_lhs) = as_u128(const_lhs) {
+                                if let Some($fold_int_rhs) = as_u128(const_rhs) {
+                                    return self.const_uint_big(result_type, $fold_int);
+                                }
+                            }
+                        })?
+                        _ => {}
+                    }
+                }
+            })?
+
             self.emit()
                 .$inst_name(result_type, None, lhs.def(self), rhs.def(self))
                 .unwrap()
@@ -600,7 +634,14 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
     simple_op! {sub, i_sub}
     simple_op! {fsub, f_sub}
     simple_op! {fsub_fast, f_sub} // fast=normal
-    simple_op! {mul, i_mul}
+    simple_op! {
+        mul, i_mul,
+        // HACK(eddyb) `rustc_codegen_ssa` relies on `Builder` methods doing
+        // on-the-fly constant-folding, for e.g. intrinsics that copy memory.
+        fold_const {
+            int(a, b) => a * b
+        }
+    }
     simple_op! {fmul, f_mul}
     simple_op! {fmul_fast, f_mul} // fast=normal
     simple_op! {udiv, u_div}
diff --git a/crates/spirv-builder/src/test/basic.rs b/crates/spirv-builder/src/test/basic.rs
@@ -440,7 +440,6 @@ pub fn main() {
 }
 
 #[test]
-#[ignore]
 fn create_uninitialized_memory() {
     val(r#"
 use core::mem::MaybeUninit;
@@ -598,3 +597,129 @@ pub fn main(image: UniformConstant<Image2d>, mut output: Output<glam::Vec4>) {
 }
 "#);
 }
+
+/// Helper to generate all of the `ptr_*` tests below, which test that the various
+/// ways to use raw pointer `read`/`write`/`copy`, to copy a single value, work,
+/// and that the resulting SPIR-V uses either a pair of `OpLoad` and `OpStore`,
+/// and/or the `OpCopyMemory` instruction, but *not* `OpCopyMemorySized`.
+macro_rules! test_copy_via_raw_ptr {
+    ($copy_expr:literal => $spirv:literal) => {
+        dis_fn(
+            concat!(
+                r#"
+        fn copy_via_raw_ptr(src: &f32, dst: &mut f32) {
+            unsafe {
+                "#,
+                $copy_expr,
+                r#"
+            }
+        }
+        #[spirv(fragment)]
+        pub fn main(i: Input<f32>, mut o: Output<f32>) {
+            copy_via_raw_ptr(&i, &mut o);
+            // FIXME(eddyb) above call results in inlining `copy_via_raw_ptr`,
+            // due to the to `Input`/`Output` storage classes, so to get the
+            // disassembled function we also need `Function`-local pointers:
+            let (src, mut dst) = (0.0, 0.0);
+            copy_via_raw_ptr(&src, &mut dst);
+        }
+"#
+            ),
+            "copy_via_raw_ptr",
+            concat!(
+                r#"%1 = OpFunction %2 None %3
+                %4 = OpFunctionParameter %5
+                %6 = OpFunctionParameter %5
+                %7 = OpLabel"#,
+                $spirv,
+                r#"OpReturn
+                OpFunctionEnd"#
+            ),
+        );
+    };
+}
+
+#[test]
+fn ptr_read() {
+    test_copy_via_raw_ptr!(
+        "*dst = core::ptr::read(src)"=> r#"
+            %8 = OpVariable %5 Function
+            OpStore %8 %9
+            OpCopyMemory %8 %4
+            %10 = OpLoad %11 %8
+            OpStore %6 %10
+        "#
+    );
+}
+
+#[test]
+fn ptr_read_method() {
+    test_copy_via_raw_ptr!(
+        "*dst = (src as *const f32).read()" => r#"
+            %8 = OpVariable %5 Function
+            OpStore %8 %9
+            OpCopyMemory %8 %4
+            %10 = OpLoad %11 %8
+            OpStore %6 %10
+        "#
+    );
+}
+
+#[test]
+fn ptr_write() {
+    test_copy_via_raw_ptr!(
+        "core::ptr::write(dst, *src)" => r#"
+            %8 = OpVariable %5 Function
+            %9 = OpLoad %10 %4
+            OpStore %8 %9
+            OpCopyMemory %6 %8
+        "#
+    );
+}
+
+#[test]
+fn ptr_write_method() {
+    test_copy_via_raw_ptr!(
+        "(dst as *mut f32).write(*src)" => r#"
+            %8 = OpVariable %5 Function
+            %9 = OpLoad %10 %4
+            OpStore %8 %9
+            OpCopyMemory %6 %8
+        "#
+    );
+}
+
+#[test]
+fn ptr_copy() {
+    test_copy_via_raw_ptr!(
+        "core::ptr::copy(src, dst, 1)" => r#"
+            OpCopyMemory %6 %4
+        "#
+    );
+}
+
+#[test]
+// FIXME(eddyb) doesn't work because `<*const T>::copy_to` is a method that wraps
+// the actual `core::ptr::copy` intrinsic - this requires either MIR inlining, or
+// making the methods themselves intrinsic (via attributes instead of pseudo-ABI).
+#[ignore]
+fn ptr_copy_to_method() {
+    test_copy_via_raw_ptr!(
+        "(src as *const f32).copy_to(dst, 1)" => r#"
+            OpCopyMemory %6 %4
+        "#
+    );
+}
+
+#[test]
+// FIXME(eddyb) doesn't work because `<*mut T>::copy_from` is a method that wraps
+// the actual `core::ptr::copy` intrinsic - this requires either MIR inlining, or
+// making the methods themselves intrinsic (via attributes instead of pseudo-ABI).
+#[ignore]
+fn ptr_copy_from_method() {
+    test_copy_via_raw_ptr!(
+        "(dst as *mut f32).copy_from(src, 1)" => r#"
+            OpCopyMemory %6 %4
+        "#
+    );
+}
diff --git a/crates/spirv-builder/src/test/mod.rs b/crates/spirv-builder/src/test/mod.rs
@@ -136,7 +136,13 @@ fn dis_fn(src: &str, func: &str, expect: &str) {
             inst.class.opcode == rspirv::spirv::Op::Name
                 && inst.operands[1].unwrap_literal_string() == abs_func_path
         })
-        .expect("No function with that name found")
+        .unwrap_or_else(|| {
+            panic!(
+                "no function with the name `{}` found in:\n{}\n",
+                abs_func_path,
+                module.disassemble()
+            )
+        })
         .operands[0]
         .unwrap_id_ref();
     let mut func = module
