From 7a5948f7298ed90dd1b7471e220204d6c22a2e72 Mon Sep 17 00:00:00 2001
From: Afonso Bordado <afonsobordado@az8.co>
Date: Thu, 24 Jun 2021 16:19:25 +0100
Subject: [PATCH] aarch64: Implement lowering i128 select

---
 .../codegen/src/isa/aarch64/lower_inst.rs     | 39 +++++++++++++------
 .../filetests/isa/aarch64/condops.clif        | 11 ++++++
 .../filetests/runtests/i128-select.clif       | 21 ++++++++++
 3 files changed, 60 insertions(+), 11 deletions(-)
 create mode 100644 cranelift/filetests/filetests/runtests/i128-select.clif

diff --git a/cranelift/codegen/src/isa/aarch64/lower_inst.rs b/cranelift/codegen/src/isa/aarch64/lower_inst.rs
index c26ab7689997..927ff4b3734f 100644
--- a/cranelift/codegen/src/isa/aarch64/lower_inst.rs
+++ b/cranelift/codegen/src/isa/aarch64/lower_inst.rs
@@ -1691,20 +1691,37 @@ pub(crate) fn lower_insn_to_regs<C: LowerCtx<I = Inst>>(
             };
 
             // csel.cond rd, rn, rm
-            let rd = get_output_reg(ctx, outputs[0]).only_reg().unwrap();
-            let rn = put_input_in_reg(ctx, inputs[1], NarrowValueMode::None);
-            let rm = put_input_in_reg(ctx, inputs[2], NarrowValueMode::None);
             let ty = ctx.output_ty(insn, 0);
             let bits = ty_bits(ty);
             let is_float = ty_has_float_or_vec_representation(ty);
-            if is_float && bits == 32 {
-                ctx.emit(Inst::FpuCSel32 { cond, rd, rn, rm });
-            } else if is_float && bits == 64 {
-                ctx.emit(Inst::FpuCSel64 { cond, rd, rn, rm });
-            } else if is_float && bits == 128 {
-                ctx.emit(Inst::VecCSel { cond, rd, rn, rm });
-            } else {
-                ctx.emit(Inst::CSel { cond, rd, rn, rm });
+
+            let dst = get_output_reg(ctx, outputs[0]);
+            let lhs = put_input_in_regs(ctx, inputs[1]);
+            let rhs = put_input_in_regs(ctx, inputs[2]);
+
+            let rd = dst.regs()[0];
+            let rn = lhs.regs()[0];
+            let rm = rhs.regs()[0];
+
+            match (is_float, bits) {
+                (true, 32) => ctx.emit(Inst::FpuCSel32 { cond, rd, rn, rm }),
+                (true, 64) => ctx.emit(Inst::FpuCSel64 { cond, rd, rn, rm }),
+                (true, 128) => ctx.emit(Inst::VecCSel { cond, rd, rn, rm }),
+                (false, 128) => {
+                    ctx.emit(Inst::CSel {
+                        cond,
+                        rd: dst.regs()[0],
+                        rn: lhs.regs()[0],
+                        rm: rhs.regs()[0],
+                    });
+                    ctx.emit(Inst::CSel {
+                        cond,
+                        rd: dst.regs()[1],
+                        rn: lhs.regs()[1],
+                        rm: rhs.regs()[1],
+                    });
+                }
+                (_, _) => ctx.emit(Inst::CSel { cond, rd, rn, rm }),
             }
         }
 
diff --git a/cranelift/filetests/filetests/isa/aarch64/condops.clif b/cranelift/filetests/filetests/isa/aarch64/condops.clif
index 7bb7ef62be5c..9334d31b02ed 100644
--- a/cranelift/filetests/filetests/isa/aarch64/condops.clif
+++ b/cranelift/filetests/filetests/isa/aarch64/condops.clif
@@ -53,3 +53,14 @@ block0(v0: i32, v1: i8, v2: i8):
 
 ; check: subs wzr, w0, #42
 ; nextln: csel x0, x1, x2, eq
+
+
+function %i128_select(b1, i128, i128) -> i128 {
+block0(v0: b1, v1: i128, v2: i128):
+  v3 = select.i128 v0, v1, v2
+  return v3
+}
+
+; check: subs wzr, w0, wzr
+; nextln: csel x0, x2, x4, ne
+; nextln: csel x1, x3, x5, ne
diff --git a/cranelift/filetests/filetests/runtests/i128-select.clif b/cranelift/filetests/filetests/runtests/i128-select.clif
new file mode 100644
index 000000000000..0d2393967300
--- /dev/null
+++ b/cranelift/filetests/filetests/runtests/i128-select.clif
@@ -0,0 +1,21 @@
+test run
+target aarch64
+target x86_64 machinst
+
+function %i128_select(i8, i64, i64, i64, i64) -> i64, i64 {
+block0(v0: i8, v1: i64, v2: i64, v3: i64, v4: i64):
+    v5 = icmp_imm ne v0, 0
+
+    v6 = iconcat v1, v2
+    v7 = iconcat v3, v4
+
+    v8 = select.i128 v5, v6, v7
+
+    v9, v10 = isplit v8
+    return v9, v10
+}
+; run: %i128_select(1, 0, 0, 1, 1) == [0, 0]
+; run: %i128_select(0, 0, 0, 1, 1) == [1, 1]
+
+; run: %i128_select(1, 1, 2, 3, 4) == [1, 2]
+; run: %i128_select(0, 1, 2, 3, 4) == [3, 4]