Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 1 | #!/usr/bin/env python3 |
| 2 | # Copyright 2021 Google LLC |
| 3 | # |
| 4 | # This source code is licensed under the BSD-style license found in the |
| 5 | # LICENSE file in the root directory of this source tree. |
Zhi An Ng | 580292d | 2022-01-04 15:45:03 -0800 | [diff] [blame] | 6 | """Converts hand written assembly (.S files) to C++ files using the JIT. |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 7 | |
Zhi An Ng | 580292d | 2022-01-04 15:45:03 -0800 | [diff] [blame] | 8 | Takes a single argument, an assembly file, and prints converted output to stdout. |
| 9 | """ |
| 10 | |
| 11 | import argparse |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 12 | import datetime |
| 13 | import re |
| 14 | import sys |
| 15 | |
| 16 | SPACES = r'\s*' |
| 17 | COMMA = r',' + SPACES |
| 18 | COMMENTS = SPACES + '((//\s+.+)|)$' |
| 19 | WB = r'!' |
| 20 | |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 21 | REG_NO_GROUP = r'r\d+|s\d+|d\d+|q\d+|sp|lr|pc|x\d+|(?:v\d+\.(?:\d+)?(?:d|s|h|b))' |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 22 | REG = r'(' + REG_NO_GROUP + ')' |
| 23 | IMM_NO_GROUP = r'\d+' |
| 24 | IMM = r'(' + IMM_NO_GROUP + ')' |
| 25 | REG_LANE_NO_GROUP = r'(?:' + REG_NO_GROUP + r')\[' + IMM_NO_GROUP + r'\]' |
| 26 | REG_OR_IMM = r'(' + REG_LANE_NO_GROUP + '|' + REG_NO_GROUP + '|' + IMM_NO_GROUP + ')' |
| 27 | |
| 28 | REGLIST_CONSEC = r'\{(\w+)-(\w+)\}' + SPACES |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 29 | REGLIST_INDIV = r'\{([\w.]+(?:,\s+[\w.]+)*)\}' + SPACES |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 30 | REGLIST_INDIV_REPLICATE = r'\{(\w+(?:\[\])(,\s*\w+(?:\[\]))*)\}' + SPACES |
| 31 | REGLIST_INDEX = r'\{(' + REG_LANE_NO_GROUP + ')\}' + SPACES |
| 32 | |
| 33 | APSR = 'APSR_nzcv' |
| 34 | FPSCR = '(FPSCR)' |
| 35 | |
| 36 | MEMOP = r'\[' + SPACES + REG + '\]' + SPACES |
| 37 | MEMOP_MAYBE_WB = r'\[' + SPACES + REG + '\]' + f'({WB})?' |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 38 | MEMOP_OFFSET = r'\[' + REG + COMMA + '(-?\d+)\]' + SPACES |
| 39 | MEMOP_OFFSET_MAYBE_WB = r'\[' + REG + COMMA + '(-?\d+)\]' + f'({WB})?' + SPACES |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 40 | |
| 41 | B_IMM = r'(\d+)(f|b)' |
| 42 | |
| 43 | INSTR = SPACES + r'([A-Z0-9.]+)' + SPACES |
| 44 | |
| 45 | # e.g. #ifndef __APPLE__ |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 46 | IFDEF_RE = re.compile(r'\s*#(ifndef|endif|ifdef).*') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 47 | # e.g. # Push 96 bytes |
| 48 | COMMENT_RE = re.compile(SPACES + r'((//|#)\s*.+)') |
| 49 | # e.g. 0: |
| 50 | LABEL = re.compile(r'(\w+):') |
| 51 | # e.g. NOP |
| 52 | INSTR_RE = re.compile(INSTR + COMMENTS) |
| 53 | # e.g. VPUSH {d8-d15} |
| 54 | INSTR_REGLIST_CONSEC_RE = re.compile(INSTR + REGLIST_CONSEC + COMMENTS) |
| 55 | # e.g. PUSH {r4, r5} |
| 56 | INSTR_REGLIST_LIST_RE = re.compile(INSTR + REGLIST_INDIV + COMMENTS) |
| 57 | # e.g. BX lr |
| 58 | INSTR_OP_RE = re.compile(INSTR + REG + COMMENTS) |
| 59 | # e.g. BLO 2f |
| 60 | INSTR_B_IMM = re.compile(INSTR + B_IMM + COMMENTS) |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 61 | # e.g. TBNZ x0, 4, 5f |
| 62 | INSTR_B_REG_IMM_IMM = re.compile(INSTR + REG + COMMA + IMM + COMMA + B_IMM + COMMENTS) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 63 | # e.g. .p2align 3 |
| 64 | P2ALIGN_RE = re.compile(SPACES + r'\.p2align\s+(\d+)') |
| 65 | # e.g. CMP r0, 2 |
| 66 | INSTR_REG_IMM_RE = re.compile(INSTR + REG + COMMA + IMM + COMMENTS) |
| 67 | # e.g. LDR r0, [r12] |
| 68 | INSTR_REG_MEMOP_RE = re.compile(INSTR + REG + COMMA + MEMOP + COMMENTS) |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 69 | # e.g. LDR q0, [x4], 16 |
| 70 | INSTR_REG_MEMOP_IMM_RE = re.compile(INSTR + REG + COMMA + MEMOP + COMMA + IMM + COMMENTS) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 71 | # e.g. LDR r0, [sp, 112] |
| 72 | INSTR_REG_MEMOP_OFFSET_RE = re.compile(INSTR + REG + COMMA + MEMOP_OFFSET + |
| 73 | COMMENTS) |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 74 | # e.g. LDRD r6, r7, [sp] |
| 75 | INSTR_REG_REG_MEMOP_RE = re.compile(INSTR + REG + COMMA + REG + COMMA + |
| 76 | MEMOP + COMMENTS) |
| 77 | # e.g. LDRD r6, r7, [sp, 104], STP d8, d9, [sp, -64]! |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 78 | INSTR_REG_REG_MEMOP_OFFSET_RE = re.compile(INSTR + REG + COMMA + REG + COMMA + |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 79 | MEMOP_OFFSET_MAYBE_WB + COMMENTS) |
| 80 | # e.g. LDP q20, q21, [x5], 32 |
| 81 | INSTR_REG_REG_MEMOP_IMM_RE = re.compile(INSTR + REG + COMMA + REG + COMMA + |
| 82 | MEMOP + COMMA + IMM + COMMENTS) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 83 | # e.g. PLD [r4, 64] |
| 84 | INSTR_MEMOP_OFFSET_RE = re.compile(INSTR + MEMOP_OFFSET + COMMENTS) |
| 85 | # e.g. movlo r12, r3, vdup.32 q0, d14[0] |
| 86 | INSTR_REG_REG_RE = re.compile(INSTR + REG + COMMA + REG_OR_IMM + COMMENTS) |
| 87 | # e.g. SUBS r5, r2, 16 or SUBS r5, r2, r10 or VMLFA.F32 q8, q4, d0[0] |
| 88 | INSTR_REG_REG_REG_RE = re.compile(INSTR + REG + COMMA + REG + COMMA + |
| 89 | REG_OR_IMM + COMMENTS) |
| 90 | # e.g. VEXT.8 q0, q0, q0, 4 |
| 91 | INSTR_REG_REG_REG_IMM_RE = re.compile(INSTR + REG + COMMA + REG + COMMA + REG + |
| 92 | COMMA + IMM + COMMENTS) |
| 93 | # e.g. VST1.32 {d16}, [r11], r0 |
| 94 | INSTR_REGLIST_INDIV_MEMOP_REG = re.compile(INSTR + REGLIST_INDIV + COMMA + |
| 95 | MEMOP + COMMA + REG + COMMENTS) |
| 96 | # e.g. VST1.32 {d16-d19}, [r11], r0 |
| 97 | INSTR_REGLIST_CONSEC_MEMOP_REG = re.compile(INSTR + REGLIST_CONSEC + COMMA + |
| 98 | MEMOP + COMMA + REG + COMMENTS) |
| 99 | # e.g. VLDM r9, {d16-d19} |
| 100 | INSTR_REG_REGLIST_CONSECT = re.compile(INSTR + REG + COMMA + REGLIST_CONSEC + |
| 101 | COMMENTS) |
| 102 | # e.g. VLDM r9!, {d16-d19} |
| 103 | INSTR_REG_REGLIST_CONSECT_WB = re.compile(INSTR + REG + WB + COMMA + |
| 104 | REGLIST_CONSEC + COMMENTS) |
| 105 | # e.g. VLDM r9!, {d16} |
| 106 | INSTR_REG_REGLIST_INDIV_WB = re.compile(INSTR + REG + WB + COMMA + |
| 107 | REGLIST_INDIV + COMMENTS) |
| 108 | # e.g. VLD1.32 {d0}, [r3]{!} |
| 109 | INSTR_REGLIST_INDIV_MEMOP = re.compile(INSTR + REGLIST_INDIV + COMMA + |
| 110 | MEMOP_MAYBE_WB + COMMENTS) |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 111 | # e.g. LD1 {v16.16b, v17.16b, v18.16b}, [x5], 48 |
| 112 | INSTR_REGLIST_INDIV_MEMOP_IMM = re.compile(INSTR + REGLIST_INDIV + COMMA + |
| 113 | MEMOP + COMMA + IMM + COMMENTS) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 114 | # e.g. VST1.32 {d24-d25}, [r11]{!} |
| 115 | INSTR_REGLIST_CONSEC_MEMOP = re.compile(INSTR + REGLIST_CONSEC + COMMA + |
| 116 | MEMOP_MAYBE_WB + COMMENTS) |
| 117 | # e.g. VLD1.32 {d0[]}, [r3]! |
| 118 | INSTR_REGLIST_REPLICATE_MEMOP = re.compile(INSTR + REGLIST_INDIV_REPLICATE + |
| 119 | COMMA + MEMOP + r'(!)?' + COMMENTS) |
| 120 | # e.g. VST1.32 {d16[0]}, [r11]{!} |
| 121 | INSTR_REGLIST_INDEX_MEMOP = re.compile(INSTR + REGLIST_INDEX + COMMA + |
| 122 | MEMOP_MAYBE_WB + COMMENTS) |
| 123 | # e.g. VMRS APSR_nzcv, FPSCR |
| 124 | INSTR_REG_FPSCR = re.compile(INSTR + f'({APSR}|{REG_NO_GROUP})' + COMMA + |
| 125 | FPSCR + COMMENTS) |
| 126 | |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 127 | # e.g. PRFM PLDL1KEEP, [x5] |
| 128 | INSTR_PLD_MEMOP = re.compile(INSTR + f'(PLDL1KEEP)' + COMMA + MEMOP + COMMENTS) |
| 129 | # e.g. PRFM PLDL1KEEP, [x5, 64] |
| 130 | INSTR_PLD_MEMOP_OFFSET = re.compile(INSTR + f'(PLDL1KEEP)' + COMMA + MEMOP_OFFSET + COMMENTS) |
| 131 | |
| 132 | COND = r'([A-Z]+)' |
| 133 | # e.g. CSEL x9, x3, x9, LO |
| 134 | INSTR_REG_REG_REG_COND_RE = re.compile(INSTR + REG + COMMA + REG + COMMA + REG + COMMA + COND + COMMENTS) |
| 135 | |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 136 | |
| 137 | def remove_brackets(s): |
| 138 | return s.replace('[', '').replace(']', '') |
| 139 | |
| 140 | |
| 141 | def fix_replicate_instruction(s): |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 142 | return re.sub(r'_(\d+)', r'r_\1', s, 1) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 143 | |
| 144 | |
| 145 | def fix_instr_name(s): |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 146 | return s.lower().replace('.', '_', 2).replace('and', 'and_', 1) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 147 | |
| 148 | |
| 149 | def fix_comments(s): |
| 150 | return s.replace('#', '//', 1) |
| 151 | |
| 152 | |
| 153 | def maybe_wb(wb): |
| 154 | return '++' if wb else '' |
| 155 | |
| 156 | |
| 157 | def fix_fn_name(name): |
| 158 | if name.startswith('xnn_'): |
| 159 | name = name[len('xnn_'):] |
| 160 | # remove any type of activations from name |
| 161 | if 'minmax' in name: |
| 162 | name = name.replace('minmax_', '') |
| 163 | return f'xnn_generate_{name}' |
| 164 | |
| 165 | |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 166 | def fix_regs(regs): |
| 167 | # Vector registers with datatype need to be method calls. |
| 168 | # e.g. v2.4s -> v2.v4s(), v2.s -> v2.s() |
| 169 | def repl(m): |
| 170 | if m.group(2): |
| 171 | return f'{m[1]}v{m[2]}{m[3]}()' |
| 172 | else: |
| 173 | return f'{m[1]}{m[3]}()' |
| 174 | return re.sub(r'(\w+\.)(\d+)?(\w+)', repl, regs) |
| 175 | |
| 176 | |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 177 | IGNORE_LINES = [r'\s*\.\w+'] |
| 178 | |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 179 | AARCH32 = 'aarch32' |
| 180 | AARCH64 = 'aarch64' |
| 181 | GEMM = 'GEMM' |
| 182 | IGEMM = 'IGEMM' |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 183 | |
Zhi An Ng | 580292d | 2022-01-04 15:45:03 -0800 | [diff] [blame] | 184 | def main(input_file): |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 185 | arch = None |
| 186 | kernel_type = GEMM |
| 187 | |
| 188 | if 'aarch32' in input_file: |
| 189 | arch = AARCH32 |
| 190 | elif 'aarch64' in input_file: |
| 191 | arch = AARCH64 |
| 192 | else: |
| 193 | print('ERROR: unknown architecture') |
| 194 | sys.exit(1) |
| 195 | |
| 196 | if 'igemm' in input_file: |
| 197 | kernel_type = IGEMM |
| 198 | |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 199 | # Whether we are in the copyright section. |
| 200 | in_copyright = False |
| 201 | # Whether we are in the microkernel function. |
| 202 | in_function = False |
| 203 | # Instructions that make up the microkernel. |
| 204 | instructions = [] |
| 205 | # Lines of code or comments before the actual function body. |
| 206 | prologue = [] |
| 207 | # All labels need to be declared first, collect them and output them after |
| 208 | # function signature. |
| 209 | labels = [] |
| 210 | # Name of the microkernel function. |
| 211 | fn_name = '' |
| 212 | sc = ';' |
Zhi An Ng | 13b57dd | 2022-01-06 09:33:20 -0800 | [diff] [blame] | 213 | # Whether we are in the auto-generated comment. |
| 214 | in_autogen = False |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 215 | |
| 216 | with open(input_file, 'r', encoding='utf-8') as f: |
| 217 | for line in f: |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 218 | line = line.rstrip() |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 219 | |
| 220 | # Handle all lines before the microkernel instructions begin. |
| 221 | if not in_function: |
Zhi An Ng | 13b57dd | 2022-01-06 09:33:20 -0800 | [diff] [blame] | 222 | if 'Auto-generated file' in line: |
| 223 | in_autogen = True |
| 224 | continue |
| 225 | elif 'BEGIN_FUNCTION' in line: |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 226 | in_function = True |
| 227 | fn_name = line.split()[1] |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 228 | prologue.append(f'// Converted from: {input_file[20:]}') |
| 229 | if kernel_type == GEMM: |
| 230 | prologue.append('void Generator::generate(size_t nc, size_t kc, void* params) {') |
| 231 | else: |
| 232 | prologue.append('void Generator::generate(size_t nc, size_t kc, size_t ks, void* params) {') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 233 | continue |
| 234 | elif 'Copyright ' in line: |
Zhi An Ng | 13b57dd | 2022-01-06 09:33:20 -0800 | [diff] [blame] | 235 | in_autogen = False |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 236 | # replace year |
| 237 | prologue.append( |
| 238 | re.sub('\d{4}', str(datetime.date.today().year), line, |
| 239 | 1).rstrip()) |
| 240 | continue |
| 241 | elif '#include <xnnpack/assembly.h>' in line: |
Zhi An Ng | 6b72e6c | 2022-02-03 11:16:27 -0800 | [diff] [blame^] | 242 | prologue.append(f'#include <cstddef>') |
| 243 | prologue.append('') |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 244 | prologue.append(f'#include <xnnpack/{arch}-assembler.h>') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 245 | prologue.append('#include <xnnpack/allocator.h>') |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 246 | if kernel_type == GEMM: |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 247 | prologue.append('#include <xnnpack/gemm.h>') |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 248 | else: |
| 249 | prologue.append('#include <xnnpack/igemm.h>') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 250 | prologue.append('') |
| 251 | prologue.append('namespace xnnpack {') |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 252 | prologue.append(f'namespace {arch} {{') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 253 | prologue.append('namespace {') |
| 254 | prologue.append('class Generator : public Assembler {') |
| 255 | prologue.append(' using Assembler::Assembler;') |
| 256 | prologue.append(' public:') |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 257 | if kernel_type == GEMM: |
| 258 | prologue.append(' void generate(size_t nc, size_t kc, void* params);') |
| 259 | else: |
| 260 | prologue.append(' void generate(size_t nc, size_t kc, size_t ks, void* params);') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 261 | prologue.append('};') |
| 262 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 263 | elif any(re.fullmatch(p, line) for p in IGNORE_LINES): |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 264 | continue |
Zhi An Ng | 13b57dd | 2022-01-06 09:33:20 -0800 | [diff] [blame] | 265 | elif in_autogen: |
| 266 | continue |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 267 | else: |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 268 | prologue.append(fix_comments(line.rstrip())) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 269 | continue |
| 270 | |
| 271 | # We are now in the microkernel function body. |
| 272 | # Don't keep the ifdefs. |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 273 | m = re.fullmatch(IFDEF_RE, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 274 | if m: |
| 275 | continue |
| 276 | # But keep other comments. |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 277 | m = re.fullmatch(COMMENT_RE, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 278 | if m: |
| 279 | instructions.append(m[1]) |
| 280 | continue |
| 281 | |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 282 | m = re.fullmatch(LABEL, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 283 | if m: |
| 284 | labels.append(m[1]) |
| 285 | instructions.append(f'bind(l{m[1]}){sc}') |
| 286 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 287 | m = re.fullmatch(INSTR_RE, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 288 | if m: |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 289 | instructions.append(f'{fix_instr_name(m[1])}(){sc} {m[2]}') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 290 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 291 | m = re.fullmatch(INSTR_OP_RE, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 292 | if m: |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 293 | instructions.append(f'{fix_instr_name(m[1])}({m[2]}){sc} {m[3]}') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 294 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 295 | m = re.fullmatch(INSTR_REGLIST_CONSEC_MEMOP_REG, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 296 | if m: |
| 297 | instructions.append( |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 298 | f'{fix_instr_name(m[1])}({{{m[2]}-{m[3]}}}, mem[{m[4]}], {m[5]}){sc} {m[6]}' |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 299 | ) |
| 300 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 301 | m = re.fullmatch(INSTR_REGLIST_INDIV_MEMOP_REG, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 302 | if m: |
| 303 | instructions.append( |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 304 | f'{fix_instr_name(m[1])}({{{fix_regs(m[2])}}}, mem[{m[3]}], {m[4]}){sc} {m[5]}') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 305 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 306 | m = re.fullmatch(INSTR_REGLIST_CONSEC_RE, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 307 | if m: |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 308 | instructions.append(f'{fix_instr_name(m[1])}({{{m[2]}-{m[3]}}}){sc} {m[4]}') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 309 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 310 | m = re.fullmatch(INSTR_REGLIST_LIST_RE, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 311 | if m: |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 312 | instructions.append(f'{fix_instr_name(m[1])}({{{m[2]}}}){sc} {m[3]}') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 313 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 314 | m = re.fullmatch(INSTR_MEMOP_OFFSET_RE, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 315 | if m: |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 316 | instructions.append(f'{fix_instr_name(m[1])}(mem[{m[2]}, {m[3]}]){sc} {m[4]}') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 317 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 318 | m = re.fullmatch(INSTR_REG_MEMOP_RE, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 319 | if m: |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 320 | instructions.append(f'{fix_instr_name(m[1])}({m[2]}, mem[{m[3]}]){sc} {m[4]}') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 321 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 322 | m = re.fullmatch(INSTR_REG_MEMOP_IMM_RE , line) |
| 323 | if m: |
| 324 | instructions.append(f'{fix_instr_name(m[1])}({m[2]}, mem[{m[3]}], {m[4]}){sc} {m[5]}') |
| 325 | continue |
| 326 | m = re.fullmatch(INSTR_REG_MEMOP_OFFSET_RE, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 327 | if m: |
| 328 | instructions.append( |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 329 | f'{fix_instr_name(m[1])}({m[2]}, mem[{m[3]}, {m[4]}]){sc} {m[5]}') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 330 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 331 | m = re.fullmatch(INSTR_REG_REG_MEMOP_RE, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 332 | if m: |
| 333 | instructions.append( |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 334 | f'{fix_instr_name(m[1])}({m[2]}, {m[3]}, mem[{m[4]}]){sc} {m[5]}') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 335 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 336 | m = re.fullmatch(INSTR_REG_REG_MEMOP_OFFSET_RE, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 337 | if m: |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 338 | if m[6]: # wb |
| 339 | instructions.append( |
| 340 | f'{fix_instr_name(m[1])}({m[2]}, {m[3]}, mem[{m[4]}, {m[5]}]++){sc} {m[7]}') |
| 341 | else: #no wb |
| 342 | instructions.append( |
| 343 | f'{fix_instr_name(m[1])}({m[2]}, {m[3]}, mem[{m[4]}, {m[5]}]){sc} {m[7]}') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 344 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 345 | m = re.fullmatch(INSTR_REG_REG_MEMOP_IMM_RE , line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 346 | if m: |
| 347 | instructions.append( |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 348 | f'{fix_instr_name(m[1])}({m[2]}, {m[3]}, mem[{m[4]}], {m[5]}){sc} {m[6]}') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 349 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 350 | m = re.fullmatch(INSTR_REG_IMM_RE, line) |
| 351 | if m: |
| 352 | instructions.append(f'{fix_instr_name(m[1])}({fix_regs(m[2])}, {m[3]}){sc} {m[4]}') |
| 353 | continue |
| 354 | m = re.fullmatch(INSTR_REG_REG_REG_RE, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 355 | if m: |
| 356 | instructions.append( |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 357 | f'{fix_instr_name(m[1])}({fix_regs(m[2])}, {fix_regs(m[3])}, {fix_regs(m[4])}){sc} {m[5]}') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 358 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 359 | m = re.fullmatch(INSTR_REG_REG_REG_IMM_RE, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 360 | if m: |
| 361 | instructions.append( |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 362 | f'{fix_instr_name(m[1])}({m[2]}, {m[3]}, {m[4]}, {m[5]}){sc} {m[6]}') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 363 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 364 | m = re.fullmatch(INSTR_REG_REG_RE, line) |
| 365 | if m: |
| 366 | instructions.append(f'{fix_instr_name(m[1])}({fix_regs(m[2])}, {fix_regs(m[3])}){sc} {m[4]}') |
| 367 | continue |
| 368 | m = re.fullmatch(INSTR_REG_REGLIST_CONSECT, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 369 | if m: |
| 370 | instructions.append( |
Zhi An Ng | c607028 | 2022-01-27 15:06:20 -0800 | [diff] [blame] | 371 | f'{fix_instr_name(m[1])}(mem[{m[2]}], {{{m[3]}-{m[4]}}}){sc} {m[5]}') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 372 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 373 | m = re.fullmatch(INSTR_REG_REGLIST_CONSECT_WB, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 374 | if m: |
| 375 | instructions.append( |
Zhi An Ng | c607028 | 2022-01-27 15:06:20 -0800 | [diff] [blame] | 376 | f'{fix_instr_name(m[1])}(mem[{m[2]}]++, {{{m[3]}-{m[4]}}}){sc} {m[5]}') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 377 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 378 | m = re.fullmatch(INSTR_REG_REGLIST_INDIV_WB, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 379 | if m: |
| 380 | instructions.append( |
Zhi An Ng | c607028 | 2022-01-27 15:06:20 -0800 | [diff] [blame] | 381 | f'{fix_instr_name(m[1])}(mem[{m[2]}]++, {{{m[3]}}}){sc} {m[4]}') |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 382 | continue |
| 383 | m = re.fullmatch(INSTR_B_IMM, line) |
| 384 | if m: |
| 385 | instructions.append(f'{fix_instr_name(m[1])}(l{m[2]}){sc} {m[4]}') |
| 386 | continue |
| 387 | m = re.fullmatch(INSTR_B_REG_IMM_IMM , line) |
| 388 | if m: |
| 389 | instructions.append(f'{fix_instr_name(m[1])}({m[2]}, {m[3]}, l{m[4]}){sc} {m[6]}') |
| 390 | continue |
| 391 | m = re.fullmatch(INSTR_REGLIST_INDIV_MEMOP, line) |
| 392 | if m: |
| 393 | instructions.append( |
| 394 | f'{fix_instr_name(m[1])}({{{fix_regs(m[2])}}}, mem[{m[3]}]{maybe_wb(m[4])}){sc} {m[5]}' |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 395 | ) |
| 396 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 397 | m = re.fullmatch(INSTR_REGLIST_INDIV_MEMOP_IMM, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 398 | if m: |
| 399 | instructions.append( |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 400 | f'{fix_instr_name(m[1])}({{{fix_regs(m[2])}}}, mem[{m[3]}], {m[4]}){sc} {m[5]}' |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 401 | ) |
| 402 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 403 | m = re.fullmatch(INSTR_REGLIST_CONSEC_MEMOP, line) |
| 404 | if m: |
| 405 | instructions.append( |
| 406 | f'{fix_instr_name(m[1])}({{{m[2]}-{m[3]}}}, mem[{m[4]}]{maybe_wb(m[5])}){sc} {m[6]}' |
| 407 | ) |
| 408 | continue |
| 409 | m = re.fullmatch(INSTR_REGLIST_REPLICATE_MEMOP, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 410 | if m: |
| 411 | if m[5]: |
| 412 | instructions.append( |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 413 | f'{fix_replicate_instruction(fix_instr_name(m[1]))}({{{remove_brackets(m[2])}}}, mem[{m[4]}]++){sc} {m[6]}' |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 414 | ) |
| 415 | else: |
| 416 | instructions.append( |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 417 | f'{fix_replicate_instruction(fix_instr_name(m[1]))}({{{remove_brackets(m[2])}}}, mem[{m[4]}]){sc} {m[6]}' |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 418 | ) |
| 419 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 420 | m = re.fullmatch(INSTR_REGLIST_INDEX_MEMOP, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 421 | if m: |
| 422 | instructions.append( |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 423 | f'{fix_instr_name(m[1])}({{{m[2]}}}, mem[{m[3]}]{maybe_wb(m[4])}){sc} {m[5]}' |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 424 | ) |
| 425 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 426 | m = re.fullmatch(P2ALIGN_RE, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 427 | if m: |
| 428 | instructions.append(f'align({1 << int(m[1])}){sc}') |
| 429 | continue |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 430 | m = re.fullmatch(INSTR_REG_FPSCR, line) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 431 | if m: |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 432 | instructions.append(f'{fix_instr_name(m[1])}({m[2]}, {m[3]}){sc} {m[4]}') |
| 433 | continue |
| 434 | m = re.fullmatch(INSTR_PLD_MEMOP, line) |
| 435 | if m: |
Zhi An Ng | 6b72e6c | 2022-02-03 11:16:27 -0800 | [diff] [blame^] | 436 | instructions.append(f'{fix_instr_name(m[1])}(k{m[2]}, mem[{m[3]}]){sc} {m[4]}') |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 437 | continue |
| 438 | m = re.fullmatch(INSTR_PLD_MEMOP_OFFSET, line) |
| 439 | if m: |
Zhi An Ng | 6b72e6c | 2022-02-03 11:16:27 -0800 | [diff] [blame^] | 440 | instructions.append(f'{fix_instr_name(m[1])}(k{m[2]}, mem[{m[3]}, {m[4]}]){sc} {m[5]}') |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 441 | continue |
| 442 | m = re.fullmatch(INSTR_REG_REG_REG_COND_RE, line) |
| 443 | if m: |
| 444 | instructions.append(f'{fix_instr_name(m[1])}({m[2]}, {m[3]}, {m[4]}, k{m[5]}){sc} {m[6]}') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 445 | continue |
| 446 | |
| 447 | # Keep empty lines for formatting |
| 448 | if line.strip() == '': |
| 449 | instructions.append('') |
| 450 | continue |
| 451 | |
| 452 | # Assembly directives that we don't are about. |
| 453 | if line.strip().startswith('.'): |
| 454 | continue |
| 455 | |
| 456 | if line.startswith('END_FUNCTION'): |
| 457 | continue |
| 458 | |
| 459 | # All other lines are error. |
| 460 | print(f'ERROR: {line}', file=sys.stderr) |
| 461 | sys.exit(1) |
| 462 | |
| 463 | # Actually emit the JIT codegen (to stdout). |
| 464 | for p in prologue: |
| 465 | print(p) |
| 466 | |
| 467 | labels_str = ', '.join(f'l{l}' for l in labels) |
| 468 | print(f' Label {labels_str};') |
| 469 | print() |
| 470 | |
| 471 | indent = ' ' |
| 472 | for i in instructions: |
| 473 | if i.strip().startswith('#'): |
| 474 | print(indent + fix_comments(i)) |
| 475 | elif i.strip().startswith('//'): |
| 476 | print(indent + i) |
| 477 | elif i.strip() == '': |
| 478 | print() |
| 479 | else: |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 480 | print(indent + (i).rstrip()) |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 481 | |
| 482 | print('}') |
| 483 | print('} // namespace') |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 484 | print(f'}} // {arch}') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 485 | print('} // xnnpack') |
| 486 | print('') |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 487 | if kernel_type == GEMM: |
Zhi An Ng | 6b72e6c | 2022-02-03 11:16:27 -0800 | [diff] [blame^] | 488 | print(f'xnn_status {fix_fn_name(fn_name)}(xnn_code_buffer* code, size_t nc, size_t kc, const void* params) {{') |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 489 | else: |
Zhi An Ng | 6b72e6c | 2022-02-03 11:16:27 -0800 | [diff] [blame^] | 490 | print(f'xnn_status {fix_fn_name(fn_name)}(xnn_code_buffer* code, size_t nc, size_t kc, size_t ks, const void* params) {{') |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 491 | print(f' using namespace xnnpack::{arch};') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 492 | print(' Generator g(code);') |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 493 | if kernel_type == GEMM: |
| 494 | print(' g.generate(nc, kc, nullptr);') |
| 495 | else: |
| 496 | print(' g.generate(nc, kc, ks, nullptr);') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 497 | print(' g.finalize();') |
Zhi An Ng | c2e2da8 | 2022-01-25 16:51:58 -0800 | [diff] [blame] | 498 | print(' if (g.error() != xnnpack::Error::kNoError) {') |
Zhi An Ng | b43b47a | 2021-12-23 16:27:22 -0800 | [diff] [blame] | 499 | print(' return xnn_status_invalid_state;') |
| 500 | print(' }') |
| 501 | print(' return xnn_status_success;') |
| 502 | print('}') |
| 503 | |
| 504 | |
| 505 | if __name__ == '__main__': |
Zhi An Ng | 580292d | 2022-01-04 15:45:03 -0800 | [diff] [blame] | 506 | parser = argparse.ArgumentParser(description='Convert assembly to to JIT C++, writes to stdout.') |
| 507 | parser.add_argument('input_file', help='Input assembly filename') |
| 508 | args = parser.parse_args() |
| 509 | main(args.input_file) |