Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 77 additions & 4 deletions src/codegen/compile_classical.ml
Original file line number Diff line number Diff line change
Expand Up @@ -4411,6 +4411,19 @@ module Blob = struct
idx env
)

let idx_nat64 env =
Func.share_code2 Func.Never env "Blob.idx_nat64" (("blob", I32Type), ("idx", I64Type)) [I32Type] (fun env get_blob get_idx ->
get_idx ^^
compile_const_64 0xFFFFFFFF00000000L ^^
G.i (Binary (Wasm.Values.I64 I64Op.And)) ^^
G.i (Test (Wasm.Values.I64 I64Op.Eqz)) ^^
E.else_trap_with env "Nat64 index exceeds 32-bit range" ^^
get_blob ^^
get_idx ^^
G.i (Convert (Wasm.Values.I32 I32Op.WrapI64)) ^^
idx env
)

let dyn_alloc_scratch env =
let (set_len, get_len) = new_local env "len" in
set_len ^^
Expand Down Expand Up @@ -4685,6 +4698,19 @@ module Arr = struct
idx env
)

let idx_nat64 env =
Func.share_code2 Func.Never env "Array.idx_nat64" (("array", I32Type), ("idx", I64Type)) [I32Type] (fun env get_array get_idx ->
get_idx ^^
compile_const_64 0xFFFFFFFF00000000L ^^
G.i (Binary (Wasm.Values.I64 I64Op.And)) ^^
G.i (Test (Wasm.Values.I64 I64Op.Eqz)) ^^
E.else_trap_with env "Nat64 index exceeds 32-bit range" ^^
get_array ^^
get_idx ^^
G.i (Convert (Wasm.Values.I32 I32Op.WrapI64)) ^^
idx env
)

let element_type env typ = match Type.promote typ with
| Type.Array element_type -> element_type
| _ -> assert false
Expand Down Expand Up @@ -11062,10 +11088,48 @@ let rec compile_lexp (env : E.t) ae lexp : G.t * SR.t * G.t =
(* Common code for a[e] as lexp and as exp.
Traps or pushes the pointer to the element on the stack
*)
and unwrap_toNat_const = function
| Const.Fun (_, Const.PrimWrapper
(Ir.NumConvTrapPrim (Type.(Nat8|Nat16|Nat32|Nat64 as pty), Type.Nat))) ->
Some pty
| _ -> None

and unwrap_toNat ae e = match e.it with
| Ir.PrimE (Ir.NumConvTrapPrim (Type.(Nat8|Nat16|Nat32|Nat64 as pty), Type.Nat), [inner]) ->
Some (pty, inner)
| Ir.PrimE (Ir.CallPrim _, [{it = Ir.VarE (_, callee); _}; inner]) ->
begin match VarEnv.lookup_var ae callee with
| Some (VarEnv.Const (_, c)) ->
Option.map (fun pty -> (pty, inner)) (unwrap_toNat_const c)
| _ -> None
end
| Ir.PrimE (Ir.CallPrim _,
[{it = Ir.PrimE (Ir.DotPrim name, [{it = Ir.VarE (_, mod_var); _}]); _}; inner]) ->
begin match VarEnv.lookup_var ae mod_var with
| Some (VarEnv.Const (_, Const.Obj fs)) ->
begin match List.assoc_opt name fs with
| Some (_, c) -> Option.map (fun pty -> (pty, inner)) (unwrap_toNat_const c)
| None -> None
end
| _ -> None
end
| _ -> None

and compile_array_index env ae e1 e2 =
compile_exp_vanilla env ae e1 ^^ (* offset to array payload *)
compile_exp_vanilla env ae e2 ^^ (* idx *)
Arr.idx_bigint env
(match unwrap_toNat ae e2 with
| Some (Type.(Nat8|Nat16|Nat32 as pty), inner) ->
(* Peephole: recognize arr[NatN.toNat(x)] and elide the
conversion, compiling x directly as its native type. *)
compile_exp_as env ae (SR.UnboxedWord32 pty) inner ^^
TaggedSmallWord.lsb_adjust pty ^^
Arr.idx env
| Some (Type.Nat64, inner) ->
compile_exp_as env ae (SR.UnboxedWord64 Type.Nat64) inner ^^
Arr.idx_nat64 env
| _ ->
compile_exp_vanilla env ae e2 ^^
Arr.idx_bigint env)

and compile_prim_invocation (env : E.t) ae p es at =
(* for more concise code when all arguments and result use the same sr *)
Expand Down Expand Up @@ -11255,8 +11319,17 @@ and compile_prim_invocation (env : E.t) ae p es at =
| IdxBlobPrim, [e1; e2] ->
SR.Vanilla,
compile_exp_vanilla env ae e1 ^^ (* offset to blob payload *)
compile_exp_vanilla env ae e2 ^^ (* idx *)
Blob.idx_bigint env
(match unwrap_toNat ae e2 with
| Some (Type.(Nat8|Nat16|Nat32 as pty), inner) ->
compile_exp_as env ae (SR.UnboxedWord32 pty) inner ^^
TaggedSmallWord.lsb_adjust pty ^^
Blob.idx env
| Some (Type.Nat64, inner) ->
compile_exp_as env ae (SR.UnboxedWord64 Type.Nat64) inner ^^
Blob.idx_nat64 env
| _ ->
compile_exp_vanilla env ae e2 ^^
Blob.idx_bigint env)

| BreakPrim name, [e] ->
let d = VarEnv.get_label_depth ae name in
Expand Down
49 changes: 45 additions & 4 deletions src/codegen/compile_enhanced.ml
Original file line number Diff line number Diff line change
Expand Up @@ -11430,10 +11430,45 @@ let rec compile_lexp (env : E.t) ae lexp : G.t * SR.t * G.t =
(* Common code for a[e] as lexp and as exp.
Traps or pushes the pointer to the element on the stack
*)
and unwrap_toNat_const = function
| Const.Fun (_, _, Const.PrimWrapper
(Ir.NumConvTrapPrim (Type.(Nat8|Nat16|Nat32|Nat64 as pty), Type.Nat))) ->
Some pty
| _ -> None

and unwrap_toNat ae e = match e.it with
| Ir.PrimE (Ir.NumConvTrapPrim (Type.(Nat8|Nat16|Nat32|Nat64 as pty), Type.Nat), [inner]) ->
Some (pty, inner)
| Ir.PrimE (Ir.CallPrim _, [{it = Ir.VarE (_, callee); _}; inner]) ->
begin match VarEnv.lookup_var ae callee with
| Some (VarEnv.Const c) ->
Option.map (fun pty -> (pty, inner)) (unwrap_toNat_const c)
| _ -> None
end
| Ir.PrimE (Ir.CallPrim _,
[{it = Ir.PrimE (Ir.DotPrim name, [{it = Ir.VarE (_, mod_var); _}]); _}; inner]) ->
begin match VarEnv.lookup_var ae mod_var with
| Some (VarEnv.Const (Const.Obj fs)) ->
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess one could even iterate this for deeper projection paths, but probably not worth it.

begin match List.assoc_opt name fs with
| Some c -> Option.map (fun pty -> (pty, inner)) (unwrap_toNat_const c)
| None -> None
end
| _ -> None
end
| _ -> None

and compile_array_index env ae e1 e2 =
compile_exp_vanilla env ae e1 ^^ (* offset to array payload *)
compile_exp_vanilla env ae e2 ^^ (* idx *)
Arr.idx_bigint env
(match unwrap_toNat ae e2 with
| Some (pty, inner) ->
(* Peephole: recognize arr[NatN.toNat(x)] and elide the
conversion, compiling x directly as its native type. *)
compile_exp_as env ae (SR.UnboxedWord64 pty) inner ^^
TaggedSmallWord.lsb_adjust pty ^^
Arr.idx env
| None ->
compile_exp_vanilla env ae e2 ^^
Arr.idx_bigint env)

and compile_prim_invocation (env : E.t) ae p es at =
(* for more concise code when all arguments and result use the same sr *)
Expand Down Expand Up @@ -11624,8 +11659,14 @@ and compile_prim_invocation (env : E.t) ae p es at =
| IdxBlobPrim, [e1; e2] ->
SR.Vanilla,
compile_exp_vanilla env ae e1 ^^ (* offset to blob payload *)
compile_exp_vanilla env ae e2 ^^ (* idx *)
Blob.idx_bigint env
(match unwrap_toNat ae e2 with
| Some (pty, inner) ->
compile_exp_as env ae (SR.UnboxedWord64 pty) inner ^^
TaggedSmallWord.lsb_adjust pty ^^
Blob.idx env
| None ->
compile_exp_vanilla env ae e2 ^^
Blob.idx_bigint env)

| BreakPrim name, [e] ->
let d = VarEnv.get_label_depth ae name in
Expand Down
164 changes: 164 additions & 0 deletions test/bench/nat64-index.mo
Original file line number Diff line number Diff line change
@@ -0,0 +1,164 @@
// Benchmark: NatN.toNat() peephole vs plain Nat array indexing
Comment thread
q-uint marked this conversation as resolved.
//MOC-FLAG --package core $MOTOKO_CORE
import Nat64 "mo:core/Nat64";
import Nat32 "mo:core/Nat32";
import {
performanceCounter;
debugPrint;
rts_heap_size;
Array_init;
natToNat32;
natToNat64;
nat32ToNat;
nat64ToNat;
} = "mo:⛔";

persistent actor NatXIndex {

transient let arrSize = 256;
transient let arr : [var Nat64] = Array_init<Nat64>(arrSize, 0);

func counters() : (Int, Nat64) = (rts_heap_size(), performanceCounter(0));

public func setup() : async () {
var k = 0;
while (k < arrSize) {
arr[k] := natToNat64(k * 0x12345);
k += 1;
};
};

// Baseline: Nat loop counter, Nat index
public func natIndex() : async () {
let (m0, n0) = counters();
var outer = 0;
while (outer < 1000) {
var acc : Nat64 = 0;
var n = 0;
while (n < arrSize) {
acc +%= arr[n];
n += 1;
};
outer += 1;
};
let (m1, n1) = counters();
debugPrint("nat_index: " # debug_show (m1 - m0, n1 - n0));
};

// Prim: nat32ToNat()
public func nat32PrimIndex() : async () {
let (m0, n0) = counters();
let arrSize32 : Nat32 = natToNat32(arrSize);
var outer = 0;
while (outer < 1000) {
var acc : Nat64 = 0;
var n : Nat32 = 0;
while (n < arrSize32) {
acc +%= arr[nat32ToNat(n)];
n +%= 1;
};
outer += 1;
};
let (m1, n1) = counters();
debugPrint("nat32_prim_index: " # debug_show (m1 - m0, n1 - n0));
};

// Core lib: Nat32.toNat()
public func nat32CoreIndex() : async () {
let (m0, n0) = counters();
let arrSize32 : Nat32 = natToNat32(arrSize);
var outer = 0;
while (outer < 1000) {
var acc : Nat64 = 0;
var n : Nat32 = 0;
while (n < arrSize32) {
acc +%= arr[Nat32.toNat(n)];
n +%= 1;
};
outer += 1;
};
let (m1, n1) = counters();
debugPrint("nat32_core_index: " # debug_show (m1 - m0, n1 - n0));
};

// Method: n.toNat()
public func nat32MethodIndex() : async () {
let (m0, n0) = counters();
let arrSize32 : Nat32 = natToNat32(arrSize);
var outer = 0;
while (outer < 1000) {
var acc : Nat64 = 0;
var n : Nat32 = 0;
while (n < arrSize32) {
acc +%= arr[n.toNat()];
n +%= 1;
};
outer += 1;
};
let (m1, n1) = counters();
debugPrint("nat32_method_index: " # debug_show (m1 - m0, n1 - n0));
};

// Prim: nat64ToNat()
public func nat64PrimIndex() : async () {
let (m0, n0) = counters();
let arrSize64 : Nat64 = natToNat64(arrSize);
var outer = 0;
while (outer < 1000) {
var acc : Nat64 = 0;
var n : Nat64 = 0;
while (n < arrSize64) {
acc +%= arr[nat64ToNat(n)];
n += 1;
};
outer += 1;
};
let (m1, n1) = counters();
debugPrint("nat64_prim_index: " # debug_show (m1 - m0, n1 - n0));
};

// Core lib: Nat64.toNat()
public func nat64CoreIndex() : async () {
let (m0, n0) = counters();
let arrSize64 : Nat64 = natToNat64(arrSize);
var outer = 0;
while (outer < 1000) {
var acc : Nat64 = 0;
var n : Nat64 = 0;
while (n < arrSize64) {
acc +%= arr[Nat64.toNat(n)];
n += 1;
};
outer += 1;
};
let (m1, n1) = counters();
debugPrint("nat64_core_index: " # debug_show (m1 - m0, n1 - n0));
};

// Method: n.toNat()
public func nat64MethodIndex() : async () {
let (m0, n0) = counters();
let arrSize64 : Nat64 = natToNat64(arrSize);
var outer = 0;
while (outer < 1000) {
var acc : Nat64 = 0;
var n : Nat64 = 0;
while (n < arrSize64) {
acc +%= arr[n.toNat()];
n += 1;
};
outer += 1;
};
let (m1, n1) = counters();
debugPrint("nat64_method_index: " # debug_show (m1 - m0, n1 - n0));
};
};

//CALL ingress setup 0x4449444C0000
//CALL ingress natIndex 0x4449444C0000
//CALL ingress nat32PrimIndex 0x4449444C0000
//CALL ingress nat32CoreIndex 0x4449444C0000
//CALL ingress nat32MethodIndex 0x4449444C0000
//CALL ingress nat64PrimIndex 0x4449444C0000
//CALL ingress nat64CoreIndex 0x4449444C0000
//CALL ingress nat64MethodIndex 0x4449444C0000
17 changes: 17 additions & 0 deletions test/bench/ok/nat64-index.drun-run.ok
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
ingress Completed: Reply: 0x4449444c016c01b3c4b1f204680100010a00000000000000000101
ingress Completed: Reply: 0x4449444c0000
ingress Completed: Reply: 0x4449444c0000
debug.print: nat_index: (0, 25_687_306)
ingress Completed: Reply: 0x4449444c0000
debug.print: nat32_prim_index: (0, 10_567_324)
ingress Completed: Reply: 0x4449444c0000
debug.print: nat32_core_index: (0, 10_567_324)
ingress Completed: Reply: 0x4449444c0000
debug.print: nat32_method_index: (0, 10_567_324)
ingress Completed: Reply: 0x4449444c0000
debug.print: nat64_prim_index: (0, 17_991_325)
ingress Completed: Reply: 0x4449444c0000
debug.print: nat64_core_index: (0, 17_991_325)
ingress Completed: Reply: 0x4449444c0000
debug.print: nat64_method_index: (0, 17_991_325)
ingress Completed: Reply: 0x4449444c0000
Loading