From 27ceb4ae3719522e2eb519ba1cd67303ae2e7259 Mon Sep 17 00:00:00 2001
From: David Rubin <daviru007@icloud.com>
Date: Wed, 29 May 2024 18:10:00 -0700
Subject: [PATCH] riscv implement `@sqrt` for f32/f64

---
 src/arch/riscv64/CodeGen.zig | 66 ++++++++++++++++++++++++++++++++----
 test/behavior/eval.zig       |  1 -
 test/behavior/floatop.zig    |  1 -
 3 files changed, 59 insertions(+), 9 deletions(-)

diff --git a/src/arch/riscv64/CodeGen.zig b/src/arch/riscv64/CodeGen.zig
index d5440f38bb..6b386de4e2 100644
--- a/src/arch/riscv64/CodeGen.zig
+++ b/src/arch/riscv64/CodeGen.zig
@@ -1275,7 +1275,7 @@ fn genBody(func: *Func, body: []const Air.Inst.Index) InnerError!void {
             .round,
             .trunc_float,
             .neg,
-            => try func.airUnaryMath(inst),
+            => try func.airUnaryMath(inst, tag),
 
             .add_with_overflow => try func.airAddWithOverflow(inst),
             .sub_with_overflow => try func.airSubWithOverflow(inst),
@@ -3741,13 +3741,65 @@ fn airBitReverse(func: *Func, inst: Air.Inst.Index) !void {
     return func.finishAir(inst, result, .{ ty_op.operand, .none, .none });
 }
 
-fn airUnaryMath(func: *Func, inst: Air.Inst.Index) !void {
-    const tag = func.air.instructions.items(.tag)[@intFromEnum(inst)];
+fn airUnaryMath(func: *Func, inst: Air.Inst.Index, tag: Air.Inst.Tag) !void {
+    const zcu = func.bin_file.comp.module.?;
     const un_op = func.air.instructions.items(.data)[@intFromEnum(inst)].un_op;
-    const result: MCValue = if (func.liveness.isUnused(inst))
-        .unreach
-    else
-        return func.fail("TODO implementairUnaryMath {s} for {}", .{ @tagName(tag), func.target.cpu.arch });
+    const result: MCValue = if (func.liveness.isUnused(inst)) .unreach else result: {
+        const ty = func.typeOf(un_op);
+
+        const operand = try func.resolveInst(un_op);
+        const operand_bit_size = ty.bitSize(zcu);
+
+        if (!math.isPowerOfTwo(operand_bit_size))
+            return func.fail("TODO: airUnaryMath non-pow 2", .{});
+
+        const operand_reg, const operand_lock = try func.promoteReg(ty, operand);
+        defer if (operand_lock) |lock| func.register_manager.unlockReg(lock);
+
+        const dst_class = func.typeRegClass(ty);
+        const dst_reg, const dst_lock = try func.allocReg(dst_class);
+        defer func.register_manager.unlockReg(dst_lock);
+
+        switch (ty.zigTypeTag(zcu)) {
+            .Float => {
+                assert(dst_class == .float);
+
+                switch (operand_bit_size) {
+                    16, 80, 128 => return func.fail("TODO: airUnaryMath Float bit-size {}", .{operand_bit_size}),
+                    32, 64 => {},
+                    else => unreachable,
+                }
+
+                switch (tag) {
+                    .sqrt => {
+                        _ = try func.addInst(.{
+                            .tag = if (operand_bit_size == 64) .fsqrtd else .fsqrts,
+                            .ops = .rrr,
+                            .data = .{
+                                .r_type = .{
+                                    .rd = dst_reg,
+                                    .rs1 = operand_reg,
+                                    .rs2 = .f0, // unused, spec says it's 0
+                                },
+                            },
+                        });
+                    },
+                    else => return func.fail("TODO: airUnaryMath Float {s}", .{@tagName(tag)}),
+                }
+            },
+            .Int => {
+                assert(dst_class == .int);
+
+                switch (tag) {
+                    else => return func.fail("TODO: airUnaryMath Float {s}", .{@tagName(tag)}),
+                }
+            },
+            else => return func.fail("TODO: airUnaryMath ty: {}", .{ty.fmt(zcu)}),
+        }
+
+        break :result MCValue{ .register = dst_reg };
+    };
+
     return func.finishAir(inst, result, .{ un_op, .none, .none });
 }
 
diff --git a/test/behavior/eval.zig b/test/behavior/eval.zig
index 2be355e064..3b77dbf966 100644
--- a/test/behavior/eval.zig
+++ b/test/behavior/eval.zig
@@ -446,7 +446,6 @@ test "binary math operator in partially inlined function" {
     if (builtin.zig_backend == .stage2_arm) return error.SkipZigTest; // TODO
     if (builtin.zig_backend == .stage2_aarch64) return error.SkipZigTest; // TODO
     if (builtin.zig_backend == .stage2_sparc64) return error.SkipZigTest; // TODO
-    if (builtin.zig_backend == .stage2_riscv64) return error.SkipZigTest;
 
     var s: [4]u32 = undefined;
     var b: [16]u8 = undefined;
diff --git a/test/behavior/floatop.zig b/test/behavior/floatop.zig
index 65d889776a..f644465d63 100644
--- a/test/behavior/floatop.zig
+++ b/test/behavior/floatop.zig
@@ -281,7 +281,6 @@ test "@sqrt f32/f64" {
     if (builtin.zig_backend == .stage2_aarch64) return error.SkipZigTest; // TODO
     if (builtin.zig_backend == .stage2_sparc64) return error.SkipZigTest; // TODO
     if (builtin.zig_backend == .stage2_x86_64 and builtin.target.ofmt != .elf and builtin.target.ofmt != .macho) return error.SkipZigTest;
-    if (builtin.zig_backend == .stage2_riscv64) return error.SkipZigTest;
 
     try testSqrt(f32);
     try comptime testSqrt(f32);