Override carrying_mul_add in cg_llvm

Author: scottmcm

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -340,6 +340,33 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                     self.const_i32(cache_type),
                 ])
             }
+            sym::carrying_mul_add => {
+                let uty = fn_args.type_at(0);
+                let (size, signed) = uty.int_size_and_signed(self.tcx);
+                assert!(!signed);
+
+                let wide_llty = self.type_ix(size.bits() * 2);
+                let a = self.zext(args[0].immediate(), wide_llty);
+                let b = self.zext(args[1].immediate(), wide_llty);
+                let c = self.zext(args[2].immediate(), wide_llty);
+                let d = self.zext(args[3].immediate(), wide_llty);
+
+                let wide = self.unchecked_umul(a, b);
+                let wide = self.unchecked_uadd(wide, c);
+                let wide = self.unchecked_uadd(wide, d);
+
+                let narrow_llty = self.type_ix(size.bits());
+                let low = self.trunc(wide, narrow_llty);
+                let bits_const = self.const_uint(wide_llty, size.bits());
+                let high = self.lshr(wide, bits_const);
+                let high = self.trunc(high, narrow_llty);
+
+                let pair_llty = self.type_struct(&[narrow_llty, narrow_llty], false);
+                let pair = self.const_poison(pair_llty);
+                let pair = self.insert_value(pair, low, 0);
+                let pair = self.insert_value(pair, high, 1);
+                pair
+            }
             sym::ctlz
             | sym::ctlz_nonzero
             | sym::cttz

tests/codegen/intrinsics/carrying_mul_add.rs

@@ -0,0 +1,110 @@
+//@ revisions: RAW OPT
+//@ compile-flags: -C opt-level=1
+//@[RAW] compile-flags: -C no-prepopulate-passes
+//@[OPT] min-llvm-version: 19
+
+#![crate_type = "lib"]
+#![feature(core_intrinsics)]
+#![feature(core_intrinsics_fallbacks)]
+
+// Note that LLVM seems to sometimes permute the order of arguments to mul and add,
+// so these tests don't check the arguments in the optimized revision.
+
+use std::intrinsics::{carrying_mul_add, fallback};
+
+// The fallbacks are emitted even when they're never used, but optimize out.
+
+// RAW: wide_mul_u128
+// OPT-NOT: wide_mul_u128
+
+// CHECK-LABEL: @cma_u8
+#[no_mangle]
+pub unsafe fn cma_u8(a: u8, b: u8, c: u8, d: u8) -> (u8, u8) {
+    // CHECK: [[A:%.+]] = zext i8 %a to i16
+    // CHECK: [[B:%.+]] = zext i8 %b to i16
+    // CHECK: [[C:%.+]] = zext i8 %c to i16
+    // CHECK: [[D:%.+]] = zext i8 %d to i16
+    // CHECK: [[AB:%.+]] = mul nuw i16
+    // RAW-SAME: [[A]], [[B]]
+    // CHECK: [[ABC:%.+]] = add nuw i16
+    // RAW-SAME: [[AB]], [[C]]
+    // CHECK: [[ABCD:%.+]] = add nuw i16
+    // RAW-SAME: [[ABC]], [[D]]
+    // CHECK: [[LOW:%.+]] = trunc i16 [[ABCD]] to i8
+    // CHECK: [[HIGHW:%.+]] = lshr i16 [[ABCD]], 8
+    // RAW: [[HIGH:%.+]] = trunc i16 [[HIGHW]] to i8
+    // OPT: [[HIGH:%.+]] = trunc nuw i16 [[HIGHW]] to i8
+    // CHECK: [[PAIR0:%.+]] = insertvalue { i8, i8 } poison, i8 [[LOW]], 0
+    // CHECK: [[PAIR1:%.+]] = insertvalue { i8, i8 } [[PAIR0]], i8 [[HIGH]], 1
+    // OPT: ret { i8, i8 } [[PAIR1]]
+    carrying_mul_add(a, b, c, d)
+}
+
+// CHECK-LABEL: @cma_u32
+#[no_mangle]
+pub unsafe fn cma_u32(a: u32, b: u32, c: u32, d: u32) -> (u32, u32) {
+    // CHECK: [[A:%.+]] = zext i32 %a to i64
+    // CHECK: [[B:%.+]] = zext i32 %b to i64
+    // CHECK: [[C:%.+]] = zext i32 %c to i64
+    // CHECK: [[D:%.+]] = zext i32 %d to i64
+    // CHECK: [[AB:%.+]] = mul nuw i64
+    // RAW-SAME: [[A]], [[B]]
+    // CHECK: [[ABC:%.+]] = add nuw i64
+    // RAW-SAME: [[AB]], [[C]]
+    // CHECK: [[ABCD:%.+]] = add nuw i64
+    // RAW-SAME: [[ABC]], [[D]]
+    // CHECK: [[LOW:%.+]] = trunc i64 [[ABCD]] to i32
+    // CHECK: [[HIGHW:%.+]] = lshr i64 [[ABCD]], 32
+    // RAW: [[HIGH:%.+]] = trunc i64 [[HIGHW]] to i32
+    // OPT: [[HIGH:%.+]] = trunc nuw i64 [[HIGHW]] to i32
+    // CHECK: [[PAIR0:%.+]] = insertvalue { i32, i32 } poison, i32 [[LOW]], 0
+    // CHECK: [[PAIR1:%.+]] = insertvalue { i32, i32 } [[PAIR0]], i32 [[HIGH]], 1
+    // OPT: ret { i32, i32 } [[PAIR1]]
+    carrying_mul_add(a, b, c, d)
+}
+
+// CHECK-LABEL: @cma_u128
+// CHECK-SAME: sret{{.+}}dereferenceable(32){{.+}}%_0,{{.+}}%a,{{.+}}%b,{{.+}}%c,{{.+}}%d
+#[no_mangle]
+pub unsafe fn cma_u128(a: u128, b: u128, c: u128, d: u128) -> (u128, u128) {
+    // CHECK: [[A:%.+]] = zext i128 %a to i256
+    // CHECK: [[B:%.+]] = zext i128 %b to i256
+    // CHECK: [[C:%.+]] = zext i128 %c to i256
+    // CHECK: [[D:%.+]] = zext i128 %d to i256
+    // CHECK: [[AB:%.+]] = mul nuw i256
+    // RAW-SAME: [[A]], [[B]]
+    // CHECK: [[ABC:%.+]] = add nuw i256
+    // RAW-SAME: [[AB]], [[C]]
+    // CHECK: [[ABCD:%.+]] = add nuw i256
+    // RAW-SAME: [[ABC]], [[D]]
+    // CHECK: [[LOW:%.+]] = trunc i256 [[ABCD]] to i128
+    // CHECK: [[HIGHW:%.+]] = lshr i256 [[ABCD]], 128
+    // RAW: [[HIGH:%.+]] = trunc i256 [[HIGHW]] to i128
+    // OPT: [[HIGH:%.+]] = trunc nuw i256 [[HIGHW]] to i128
+    // RAW: [[PAIR0:%.+]] = insertvalue { i128, i128 } poison, i128 [[LOW]], 0
+    // RAW: [[PAIR1:%.+]] = insertvalue { i128, i128 } [[PAIR0]], i128 [[HIGH]], 1
+    // OPT: store i128 [[LOW]], ptr %_0
+    // OPT: [[P1:%.+]] = getelementptr inbounds i8, ptr %_0, i64 16
+    // OPT: store i128 [[HIGH]], ptr [[P1]]
+    // CHECK: ret void
+    carrying_mul_add(a, b, c, d)
+}
+
+// CHECK-LABEL: @fallback_cma_u32
+#[no_mangle]
+pub unsafe fn fallback_cma_u32(a: u32, b: u32, c: u32, d: u32) -> (u32, u32) {
+    // OPT-DAG: [[A:%.+]] = zext i32 %a to i64
+    // OPT-DAG: [[B:%.+]] = zext i32 %b to i64
+    // OPT-DAG: [[AB:%.+]] = mul nuw i64
+    // OPT-DAG: [[C:%.+]] = zext i32 %c to i64
+    // OPT-DAG: [[ABC:%.+]] = add nuw i64{{.+}}[[C]]
+    // OPT-DAG: [[D:%.+]] = zext i32 %d to i64
+    // OPT-DAG: [[ABCD:%.+]] = add nuw i64{{.+}}[[D]]
+    // OPT-DAG: [[LOW:%.+]] = trunc i64 [[ABCD]] to i32
+    // OPT-DAG: [[HIGHW:%.+]] = lshr i64 [[ABCD]], 32
+    // OPT-DAG: [[HIGH:%.+]] = trunc nuw i64 [[HIGHW]] to i32
+    // OPT-DAG: [[PAIR0:%.+]] = insertvalue { i32, i32 } poison, i32 [[LOW]], 0
+    // OPT-DAG: [[PAIR1:%.+]] = insertvalue { i32, i32 } [[PAIR0]], i32 [[HIGH]], 1
+    // OPT-DAG: ret { i32, i32 } [[PAIR1]]
+    fallback::carrying_mul_add(a, b, c, d)
+}