Commit f4421c01e8a8793a7c02fce870bacbb710bb1e12 | zig

author	Jakub Konka <kubkon@jakubkonka.com> 2022-05-04 21:20:31+02:00
committer	Jakub Konka <kubkon@jakubkonka.com> 2022-05-05 21:43:36+02:00
log	f4421c01e8a8793a7c02fce870bacbb710bb1e12
tree	90d2605dd38caa885d21fcc5a94eec76445df449
parent	8715b01005c49ff99327a87264ffaa28fb3807a0

aarch64: implement mul_with_overflow for ints in range 33-64 bits incl

4 files changed, 230 insertions(+), 29 deletions(-)

src
- arch
  - aarch64
    CodeGen.zig
    Emit.zig
    Mir.zig
    bits.zig

src/arch/aarch64/CodeGen.zig+195-23

...	@@ -1294,28 +1294,29 @@ fn binOpRegister(	...	@@ -1294,28 +1294,29 @@ fn binOpRegister(
1294	};	1294	};
1295	defer self.register_manager.unfreezeRegs(&.{rhs_reg});	1295	defer self.register_manager.unfreezeRegs(&.{rhs_reg});
1296		1296
1297	const dest_reg = switch (mir_tag) {	1297	const dest_reg: Register = reg: {
1298	.cmp_shifted_register => undefined, // cmp has no destination register	1298	const dest_reg = switch (mir_tag) {
1299	.smull, .umull => blk: {	1299	.cmp_shifted_register => undefined, // cmp has no destination register
1300	// TODO can we reuse anything for smull and umull?	1300	else => if (maybe_inst) \|inst\| blk: {
1301	const raw_reg = try self.register_manager.allocReg(null);	1301	const bin_op = self.air.instructions.items(.data)[inst].bin_op;
1302	break :blk raw_reg.to64();	1302
1303	},	1303	if (lhs_is_register and self.reuseOperand(inst, bin_op.lhs, 0, lhs)) {
1304	else => if (maybe_inst) \|inst\| blk: {	1304	break :blk lhs_reg;
1305	const bin_op = self.air.instructions.items(.data)[inst].bin_op;	1305	} else if (rhs_is_register and self.reuseOperand(inst, bin_op.rhs, 1, rhs)) {
1306		1306	break :blk rhs_reg;
1307	if (lhs_is_register and self.reuseOperand(inst, bin_op.lhs, 0, lhs)) {	1307	} else {
1308	break :blk lhs_reg;	1308	const raw_reg = try self.register_manager.allocReg(inst);
1309	} else if (rhs_is_register and self.reuseOperand(inst, bin_op.rhs, 1, rhs)) {	1309	break :blk registerAlias(raw_reg, lhs_ty.abiSize(self.target.*));
1310	break :blk rhs_reg;	1310	}
1311	} else {	1311	} else blk: {
1312	const raw_reg = try self.register_manager.allocReg(inst);	1312	const raw_reg = try self.register_manager.allocReg(null);
1313	break :blk registerAlias(raw_reg, lhs_ty.abiSize(self.target.*));	1313	break :blk registerAlias(raw_reg, lhs_ty.abiSize(self.target.*));
1314	}	1314	},
1315	} else blk: {	1315	};
1316	const raw_reg = try self.register_manager.allocReg(null);	1316	break :reg switch (mir_tag) {
1317	break :blk registerAlias(raw_reg, lhs_ty.abiSize(self.target.*));	1317	.smull, .umull => dest_reg.to64(),
1318	},	1318	else => dest_reg,
		1319	};
1319	};	1320	};
1320		1321
1321	if (!lhs_is_register) try self.genSetReg(lhs_ty, lhs_reg, lhs);	1322	if (!lhs_is_register) try self.genSetReg(lhs_ty, lhs_reg, lhs);
...	@@ -1340,7 +1341,9 @@ fn binOpRegister(	...	@@ -1340,7 +1341,9 @@ fn binOpRegister(
1340	.shift = .lsl,	1341	.shift = .lsl,
1341	} },	1342	} },
1342	.mul,	1343	.mul,
		1344	.smulh,
1343	.smull,	1345	.smull,
		1346	.umulh,
1344	.umull,	1347	.umull,
1345	.lsl_register,	1348	.lsl_register,
1346	.asr_register,	1349	.asr_register,
...	@@ -1946,8 +1949,177 @@ fn airMulWithOverflow(self: *Self, inst: Air.Inst.Index) !void {	...	@@ -1946,8 +1949,177 @@ fn airMulWithOverflow(self: *Self, inst: Air.Inst.Index) !void {
1946		1949
1947	break :result MCValue{ .stack_offset = stack_offset };	1950	break :result MCValue{ .stack_offset = stack_offset };
1948	} else if (int_info.bits <= 64) {	1951	} else if (int_info.bits <= 64) {
1949	return self.fail("TODO implement mul_with_overflow for ints", .{});	1952	const stack_offset = try self.allocMem(inst, tuple_size, tuple_align);
1950	} else return self.fail("TODO implmenet mul_with_overflow for integers > u64/i64", .{});	1953
		1954	try self.spillCompareFlagsIfOccupied();
		1955	self.compare_flags_inst = null;
		1956
		1957	// TODO this should really be put in a helper similar to `binOpRegister`
		1958	const lhs_is_register = lhs == .register;
		1959	const rhs_is_register = rhs == .register;
		1960
		1961	if (lhs_is_register) self.register_manager.freezeRegs(&.{lhs.register});
		1962	if (rhs_is_register) self.register_manager.freezeRegs(&.{rhs.register});
		1963
		1964	const lhs_reg = if (lhs_is_register) lhs.register else blk: {
		1965	const raw_reg = try self.register_manager.allocReg(null);
		1966	const reg = registerAlias(raw_reg, lhs_ty.abiSize(self.target.*));
		1967	self.register_manager.freezeRegs(&.{reg});
		1968	break :blk reg;
		1969	};
		1970	defer self.register_manager.unfreezeRegs(&.{lhs_reg});
		1971
		1972	const rhs_reg = if (rhs_is_register) rhs.register else blk: {
		1973	const raw_reg = try self.register_manager.allocReg(null);
		1974	const reg = registerAlias(raw_reg, rhs_ty.abiAlignment(self.target.*));
		1975	self.register_manager.freezeRegs(&.{reg});
		1976	break :blk reg;
		1977	};
		1978	defer self.register_manager.unfreezeRegs(&.{rhs_reg});
		1979
		1980	if (!lhs_is_register) try self.genSetReg(lhs_ty, lhs_reg, lhs);
		1981	if (!rhs_is_register) try self.genSetReg(rhs_ty, rhs_reg, rhs);
		1982
		1983	// TODO reuse operands
		1984	const dest_reg = blk: {
		1985	const raw_reg = try self.register_manager.allocReg(null);
		1986	const reg = registerAlias(raw_reg, lhs_ty.abiSize(self.target.*));
		1987	self.register_manager.freezeRegs(&.{reg});
		1988	break :blk reg;
		1989	};
		1990	defer self.register_manager.unfreezeRegs(&.{dest_reg});
		1991
		1992	switch (int_info.signedness) {
		1993	.signed => {
		1994	// mul dest, lhs, rhs
		1995	_ = try self.addInst(.{
		1996	.tag = .mul,
		1997	.data = .{ .rrr = .{
		1998	.rd = dest_reg,
		1999	.rn = lhs_reg,
		2000	.rm = rhs_reg,
		2001	} },
		2002	});
		2003
		2004	const dest_high_reg = try self.register_manager.allocReg(null);
		2005	self.register_manager.freezeRegs(&.{dest_high_reg});
		2006	defer self.register_manager.unfreezeRegs(&.{dest_high_reg});
		2007
		2008	// smulh dest_high, lhs, rhs
		2009	_ = try self.addInst(.{
		2010	.tag = .smulh,
		2011	.data = .{ .rrr = .{
		2012	.rd = dest_high_reg,
		2013	.rn = lhs_reg,
		2014	.rm = rhs_reg,
		2015	} },
		2016	});
		2017
		2018	// cmp dest_high, dest, asr #63
		2019	_ = try self.addInst(.{
		2020	.tag = .cmp_shifted_register,
		2021	.data = .{ .rr_imm6_shift = .{
		2022	.rn = dest_high_reg,
		2023	.rm = dest_reg,
		2024	.imm6 = 63,
		2025	.shift = .asr,
		2026	} },
		2027	});
		2028
		2029	const shift: u6 = @intCast(u6, @as(u7, 64) - @intCast(u7, int_info.bits));
		2030	if (shift > 0) {
		2031	// lsl dest_high, dest, #shift
		2032	_ = try self.addInst(.{
		2033	.tag = .lsl_immediate,
		2034	.data = .{ .rr_shift = .{
		2035	.rd = dest_high_reg,
		2036	.rn = dest_reg,
		2037	.shift = shift,
		2038	} },
		2039	});
		2040
		2041	// cmp dest, dest_high, #shift
		2042	_ = try self.addInst(.{
		2043	.tag = .cmp_shifted_register,
		2044	.data = .{ .rr_imm6_shift = .{
		2045	.rn = dest_reg,
		2046	.rm = dest_high_reg,
		2047	.imm6 = shift,
		2048	.shift = .asr,
		2049	} },
		2050	});
		2051	}
		2052	},
		2053	.unsigned => {
		2054	const dest_high_reg = try self.register_manager.allocReg(null);
		2055	self.register_manager.freezeRegs(&.{dest_high_reg});
		2056	defer self.register_manager.unfreezeRegs(&.{dest_high_reg});
		2057
		2058	// umulh dest_high, lhs, rhs
		2059	_ = try self.addInst(.{
		2060	.tag = .umulh,
		2061	.data = .{ .rrr = .{
		2062	.rd = dest_high_reg,
		2063	.rn = lhs_reg,
		2064	.rm = rhs_reg,
		2065	} },
		2066	});
		2067
		2068	// mul dest, lhs, rhs
		2069	_ = try self.addInst(.{
		2070	.tag = .mul,
		2071	.data = .{ .rrr = .{
		2072	.rd = dest_reg,
		2073	.rn = lhs_reg,
		2074	.rm = rhs_reg,
		2075	} },
		2076	});
		2077
		2078	_ = try self.binOp(
		2079	.cmp_eq,
		2080	null,
		2081	.{ .register = dest_high_reg },
		2082	.{ .immediate = 0 },
		2083	Type.usize,
		2084	Type.usize,
		2085	);
		2086
		2087	if (int_info.bits < 64) {
		2088	// lsr dest_high, dest, #shift
		2089	_ = try self.addInst(.{
		2090	.tag = .lsr_immediate,
		2091	.data = .{ .rr_shift = .{
		2092	.rd = dest_high_reg,
		2093	.rn = dest_reg,
		2094	.shift = @intCast(u6, int_info.bits),
		2095	} },
		2096	});
		2097
		2098	_ = try self.binOp(
		2099	.cmp_eq,
		2100	null,
		2101	.{ .register = dest_high_reg },
		2102	.{ .immediate = 0 },
		2103	Type.usize,
		2104	Type.usize,
		2105	);
		2106	}
		2107	},
		2108	}
		2109
		2110	const truncated_reg = try self.register_manager.allocReg(null);
		2111	self.register_manager.freezeRegs(&.{truncated_reg});
		2112	defer self.register_manager.unfreezeRegs(&.{truncated_reg});
		2113
		2114	try self.truncRegister(dest_reg, truncated_reg, int_info.signedness, int_info.bits);
		2115
		2116	try self.genSetStack(lhs_ty, stack_offset, .{ .register = truncated_reg });
		2117	try self.genSetStack(Type.initTag(.u1), stack_offset - overflow_bit_offset, .{
		2118	.compare_flags_unsigned = .neq,
		2119	});
		2120
		2121	break :result MCValue{ .stack_offset = stack_offset };
		2122	} else return self.fail("TODO implement mul_with_overflow for integers > u64/i64", .{});
1951	},	2123	},
1952	else => unreachable,	2124	else => unreachable,
1953	}	2125	}

src/arch/aarch64/Emit.zig+12-1

 @@ -167,7 +167,9 @@ pub fn emitMir(
             .movz => try emit.mirMoveWideImmediate(inst),
             .mul => try emit.mirDataProcessing3Source(inst),
+            .smulh => try emit.mirDataProcessing3Source(inst),
             .smull => try emit.mirDataProcessing3Source(inst),
+            .umulh => try emit.mirDataProcessing3Source(inst),
             .umull => try emit.mirDataProcessing3Source(inst),
             .nop => try emit.mirNop(),
 @@ -677,7 +679,14 @@ fn mirLogicalImmediate(emit: *Emit, inst: Mir.Inst.Index) !void {
     switch (tag) {
         .eor_immediate => try emit.writeInstruction(Instruction.eorImmediate(rd, rn, imms, immr, n)),
-        .tst_immediate => try emit.writeInstruction(Instruction.tstImmediate(rn, imms, immr, n)),
+        .tst_immediate => {
+            const zr: Register = switch (rd.size()) {
+                32 => .wzr,
+                64 => .xzr,
+                else => unreachable,
+            };
+            try emit.writeInstruction(Instruction.andsImmediate(zr, rn, imms, immr, n));
+        },
         else => unreachable,
     }
 }
 @@ -1004,7 +1013,9 @@ fn mirDataProcessing3Source(emit: *Emit, inst: Mir.Inst.Index) !void {
     switch (tag) {
         .mul => try emit.writeInstruction(Instruction.mul(rrr.rd, rrr.rn, rrr.rm)),
+        .smulh => try emit.writeInstruction(Instruction.smulh(rrr.rd, rrr.rn, rrr.rm)),
         .smull => try emit.writeInstruction(Instruction.smull(rrr.rd, rrr.rn, rrr.rm)),
+        .umulh => try emit.writeInstruction(Instruction.umulh(rrr.rd, rrr.rn, rrr.rm)),
         .umull => try emit.writeInstruction(Instruction.umull(rrr.rd, rrr.rn, rrr.rm)),
         else => unreachable,
     }

src/arch/aarch64/Mir.zig+4

...	@@ -146,6 +146,8 @@ pub const Inst = struct {	...	@@ -146,6 +146,8 @@ pub const Inst = struct {
146	ret,	146	ret,
147	/// Signed bitfield extract	147	/// Signed bitfield extract
148	sbfx,	148	sbfx,
		149	/// Signed multiply high
		150	smulh,
149	/// Signed multiply long	151	/// Signed multiply long
150	smull,	152	smull,
151	/// Signed extend byte	153	/// Signed extend byte
...	@@ -188,6 +190,8 @@ pub const Inst = struct {	...	@@ -188,6 +190,8 @@ pub const Inst = struct {
188	tst_immediate,	190	tst_immediate,
189	/// Unsigned bitfield extract	191	/// Unsigned bitfield extract
190	ubfx,	192	ubfx,
		193	/// Unsigned multiply high
		194	umulh,
191	/// Unsigned multiply long	195	/// Unsigned multiply long
192	umull,	196	umull,
193	/// Unsigned extend byte	197	/// Unsigned extend byte

src/arch/aarch64/bits.zig+19-5

 @@ -1409,10 +1409,6 @@ pub const Instruction = union(enum) {
         return logicalImmediate(0b11, rd, rn, imms, immr, n);
     }
-    pub fn tstImmediate(rn: Register, imms: u6, immr: u6, n: u1) Instruction {
-        return andsImmediate(.xzr, rn, imms, immr, n);
-    }
     // Bitfield
     pub fn sbfm(rd: Register, rn: Register, immr: u6, imms: u6) Instruction {
 @@ -1589,10 +1585,20 @@ pub const Instruction = union(enum) {
         return smaddl(rd, rn, rm, .xzr);
     }
+    pub fn smulh(rd: Register, rn: Register, rm: Register) Instruction {
+        assert(rd.size() == 64);
+        return dataProcessing3Source(0b00, 0b010, 0b0, rd, rn, rm, .xzr);
+    }
     pub fn umull(rd: Register, rn: Register, rm: Register) Instruction {
         return umaddl(rd, rn, rm, .xzr);
     }
+    pub fn umulh(rd: Register, rn: Register, rm: Register) Instruction {
+        assert(rd.size() == 64);
+        return dataProcessing3Source(0b00, 0b110, 0b0, rd, rn, rm, .xzr);
+    }
     pub fn mneg(rd: Register, rn: Register, rm: Register) Instruction {
         return msub(rd, rn, rm, .xzr);
     }
 @@ -1820,9 +1826,17 @@ test "serialize instructions" {
             .expected = 0b1_00_11011_0_01_00001_0_11111_00000_00000,
         },
         .{ // tst x0, #0xffffffff00000000
-            .inst = Instruction.tstImmediate(.x0, 0b011111, 0b100000, 0b1),
+            .inst = Instruction.andsImmediate(.xzr, .x0, 0b011111, 0b100000, 0b1),
             .expected = 0b1_11_100100_1_100000_011111_00000_11111,
         },
+        .{ // umulh x0, x1, x2
+            .inst = Instruction.umulh(.x0, .x1, .x2),
+            .expected = 0b1_00_11011_1_10_00010_0_11111_00001_00000,
+        },
+        .{ // smulh x0, x1, x2
+            .inst = Instruction.smulh(.x0, .x1, .x2),
+            .expected = 0b1_00_11011_0_10_00010_0_11111_00001_00000,
+        },
     };
     for (testcases) |case| {