RISC-V peephole optimization and user-defined infix operators | Danil Usoltsev & Victoria Ostrovskaya

EML/.gitignore

@@ -1,3 +1,11 @@
 /_build
 /_coverage
 
+*.o
+*.s
+*.S
+*.ll
+*.exe
+*.a
+*.so
+*.dll

EML/lib/backend/llvm_ir/analysis.ml

@@ -43,17 +43,63 @@ let analyze (program : anf_program) =
         | AnfEval _ -> None)
       program
   in
+  let is_valid_linker_ident name =
+    String.length name > 0
+    && String.for_all
+         (fun c ->
+            (c >= 'a' && c <= 'z')
+            || (c >= 'A' && c <= 'Z')
+            || (c >= '0' && c <= '9')
+            || Char.equal c '_')
+         name
+  in
+  let mangle_operator_for_linker name =
+    "op_"
+    ^ Base.String.concat_map name ~f:(function
+      | '*' -> "_star"
+      | '+' -> "_plus"
+      | '-' -> "_minus"
+      | '/' -> "_slash"
+      | '=' -> "_eq"
+      | '<' -> "_lt"
+      | '>' -> "_gt"
+      | '!' -> "_bang"
+      | '&' -> "_amp"
+      | '|' -> "_bar"
+      | '^' -> "_hat"
+      | '@' -> "_at"
+      | '~' -> "_tilde"
+      | '?' -> "_q"
+      | '.' -> "_dot"
+      | ':' -> "_colon"
+      | '%' -> "_percent"
+      | '$' -> "_dollar"
+      | c
+        when (c >= 'a' && c <= 'z')
+             || (c >= 'A' && c <= 'Z')
+             || (c >= '0' && c <= '9')
+             || Char.equal c '_' -> String.make 1 c
+      | c -> "_u" ^ Int.to_string (Char.code c))
+  in
   let mangle_reserved name =
     if is_reserved name
     then "eml_" ^ name
     else if String.equal name "_start"
     then "eml_start"
     else name
   in
+  let asm_name_for_func func_name =
+    let base =
+      if is_valid_linker_ident func_name
+      then func_name
+      else mangle_operator_for_linker func_name
+    in
+    mangle_reserved base
+  in
   let functions, _ =
     List.fold_left
       (fun (reversed_functions, counts) (func_name, _arity, params, body, is_rec) ->
-         let base_asm_name = mangle_reserved func_name in
+         let base_asm_name = asm_name_for_func func_name in
          let duplicate_index =
            Base.Map.find counts func_name |> Option.value ~default:0
          in

EML/lib/backend/llvm_ir/generator.ml

@@ -293,6 +293,7 @@ let gen_binop_native op left_v right_v =
     let* r = untag_bool_val right_v in
     let* v = with_optional_value (or_ builder l r "or") in
     tag_bool_result v
+  | Custom _ -> fail "Custom operator must be compiled to application"
 ;;
 
 let gen_unop_native op tagged_val =

EML/lib/backend/ricsv/analysis.ml

@@ -171,13 +171,59 @@ let analyze (program : anf_program) =
         | AnfEval _ -> None)
       program
   in
+  let is_valid_linker_ident name =
+    String.length name > 0
+    && String.for_all
+         (fun c ->
+            (c >= 'a' && c <= 'z')
+            || (c >= 'A' && c <= 'Z')
+            || (c >= '0' && c <= '9')
+            || Char.equal c '_')
+         name
+  in
+  let mangle_operator_for_linker name =
+    "op_"
+    ^ Base.String.concat_map name ~f:(function
+      | '*' -> "_star"
+      | '+' -> "_plus"
+      | '-' -> "_minus"
+      | '/' -> "_slash"
+      | '=' -> "_eq"
+      | '<' -> "_lt"
+      | '>' -> "_gt"
+      | '!' -> "_bang"
+      | '&' -> "_amp"
+      | '|' -> "_bar"
+      | '^' -> "_hat"
+      | '@' -> "_at"
+      | '~' -> "_tilde"
+      | '?' -> "_q"
+      | '.' -> "_dot"
+      | ':' -> "_colon"
+      | '%' -> "_percent"
+      | '$' -> "_dollar"
+      | c
+        when (c >= 'a' && c <= 'z')
+             || (c >= 'A' && c <= 'Z')
+             || (c >= '0' && c <= '9')
+             || Char.equal c '_' -> String.make 1 c
+      | c -> "_u" ^ Int.to_string (Char.code c))
+  in
   let mangle_reserved name =
     if is_reserved name
     then "eml_" ^ name
     else if String.equal name "_start"
     then "eml_start"
     else name
   in
+  let asm_name_for_func func_name =
+    let base =
+      if is_valid_linker_ident func_name
+      then func_name
+      else mangle_operator_for_linker func_name
+    in
+    mangle_reserved base
+  in
   let functions, _ =
     List.fold_left
       (fun (reversed_functions, generated_name_counts)
@@ -189,7 +235,7 @@ let analyze (program : anf_program) =
         , slots_count
         , max_stack_args
         , max_create_tuple_array_bytes ) ->
-         let base_asm_name = mangle_reserved func_name in
+         let base_asm_name = asm_name_for_func func_name in
          let duplicate_index =
            Base.Map.find generated_name_counts func_name |> Option.value ~default:0
          in

EML/lib/backend/ricsv/architecture.ml

@@ -39,6 +39,7 @@ module Riscv_backend = struct
     | Xor of reg * reg * reg (* xor двух регистров: rd = rs1 ^ rs2 *)
     | Mul of reg * reg * reg (* умножение: rd = rs1 * rs2 *)
     | Div of reg * reg * reg (* целочисленное деление: rd = rs1 / rs2 *)
+    | Slli of reg * reg * int (* логический сдвиг влево на константу: rd = rs << imm *)
     | Srli of reg * reg * int (* логический сдвиг вправо на константу: rd = rs >>> imm *)
 
   let pp_reg ppf = function
@@ -59,6 +60,7 @@ module Riscv_backend = struct
     | Sub (rd, rs1, rs2) -> fprintf ppf "sub %a, %a, %a" pp_reg rd pp_reg rs1 pp_reg rs2
     | Mul (rd, rs1, rs2) -> fprintf ppf "mul %a, %a, %a" pp_reg rd pp_reg rs1 pp_reg rs2
     | Div (rd, rs1, rs2) -> fprintf ppf "div %a, %a, %a" pp_reg rd pp_reg rs1 pp_reg rs2
+    | Slli (rd, rs1, imm) -> fprintf ppf "slli %a, %a, %d" pp_reg rd pp_reg rs1 imm
     | Srli (rd, rs1, imm) -> fprintf ppf "srli %a, %a, %d" pp_reg rd pp_reg rs1 imm
     | Xori (rd, rs1, imm) -> fprintf ppf "xori %a, %a, %d" pp_reg rd pp_reg rs1 imm
     | Xor (rd, rs1, rs2) -> fprintf ppf "xor %a, %a, %a" pp_reg rd pp_reg rs1 pp_reg rs2
@@ -113,6 +115,7 @@ module Riscv_backend = struct
   let xor rd rs1 rs2 = [ Xor (rd, rs1, rs2) ]
   let mul rd rs1 rs2 = [ Mul (rd, rs1, rs2) ]
   let div rd rs1 rs2 = [ Div (rd, rs1, rs2) ]
+  let slli rd rs imm = [ Slli (rd, rs, imm) ]
   let srli rd rs imm = [ Srli (rd, rs, imm) ]
   let add_tag_items dst delta = [ Addi (dst, dst, delta) ]
   let arg_regs = [ a0; a1; a2; a3; a4; a5; a6; a7 ]

EML/lib/backend/ricsv/architecture.mli

@@ -35,6 +35,7 @@ module Riscv_backend : sig
     | Xor of reg * reg * reg
     | Mul of reg * reg * reg
     | Div of reg * reg * reg
+    | Slli of reg * reg * int
     | Srli of reg * reg * int
 
   val pp_reg : Format.formatter -> reg -> unit
@@ -76,6 +77,7 @@ module Riscv_backend : sig
   val xor : reg -> reg -> reg -> instr list
   val mul : reg -> reg -> reg -> instr list
   val div : reg -> reg -> reg -> instr list
+  val slli : reg -> reg -> int -> instr list
   val srli : reg -> reg -> int -> instr list
   val add_tag_items : reg -> int -> instr list
   val arg_regs : reg list

EML/lib/backend/ricsv/generator.ml

@@ -420,6 +420,8 @@ and gen_cexpr dst = function
   | ComplexImmediate imm -> gen_imm dst imm
   | ComplexUnarOper (Negative, op) -> gen_neg dst op
   | ComplexUnarOper (Not, op) -> gen_not dst op
+  | ComplexBinOper (Custom _, _, _) ->
+    fail "Custom operator must be compiled to application"
   | ComplexBinOper (op, left, right) -> gen_binop dst op left right
   | ComplexBranch (cond, then_e, else_e) -> gen_branch dst cond then_e else_e
   | ComplexField (tuple_imm, idx) -> gen_field dst tuple_imm idx
@@ -453,15 +455,19 @@ let bind_param_to_stack env i = function
   | _ -> fail "unsupported pattern"
 ;;
 
-let flush_instr_buffer ppf =
+let flush_instr_buffer ~enable_peephole ppf =
   let* state = get in
   let instruction_buffer = state.instr_buffer in
   let* () = put { state with instr_buffer = [] } in
-  let () = List.iter (fun item -> format_item ppf item) (List.rev instruction_buffer) in
+  let instructions = List.rev instruction_buffer in
+  let instructions =
+    if enable_peephole then Peephole.optimize instructions else instructions
+  in
+  let () = List.iter (fun item -> format_item ppf item) instructions in
   return ()
 ;;
 
-let gen_func ~enable_gc asm_name params body frame_sz ppf =
+let gen_func ~enable_gc ~enable_peephole asm_name params body frame_sz ppf =
   fprintf ppf "\n  .globl %s\n  .type %s, @function\n" asm_name asm_name;
   let args = List.length params in
   let params_reg, params_stack =
@@ -486,11 +492,11 @@ let gen_func ~enable_gc asm_name params body frame_sz ppf =
   let* () = spill_params_to_frame params_reg in
   let* () = gen_anf result_reg body in
   let* () = append (epilogue ~enable_gc ~is_main:(String.equal asm_name "main")) in
-  let* () = flush_instr_buffer ppf in
+  let* () = flush_instr_buffer ~enable_peephole ppf in
   return ()
 ;;
 
-let gen_program ~enable_gc ppf (analysis : analysis_result) =
+let gen_program ~enable_gc ~enable_peephole ppf (analysis : analysis_result) =
   fprintf ppf ".section .text";
   let base = Runtime.Primitives.runtime_primitive_arities in
   let arity_map =
@@ -516,7 +522,7 @@ let gen_program ~enable_gc ppf (analysis : analysis_result) =
       let* () =
         modify (fun state -> { state with current_func_index = function_index })
       in
-      gen_func ~enable_gc fn.asm_name fn.params fn.body frame_sz ppf)
+      gen_func ~enable_gc ~enable_peephole fn.asm_name fn.params fn.body frame_sz ppf)
   in
   match run comp init with
   | Ok ((), _) ->

EML/lib/backend/ricsv/generator.mli

@@ -4,6 +4,7 @@
 
 val gen_program
   :  enable_gc:bool
+  -> enable_peephole:bool
   -> Format.formatter
   -> Analysis.analysis_result
   -> (unit, string) Result.t