Fusing `Map` nodes in incremental

Learning
1

Tl;dr: Is there a way to automatically turn Incr.map ~f (Incr.map ~f:g x) into a single incremental node with the same structure as Incr.map ~f:(Fn.compose f g) x?

Context:

I’m working on an incremental evaluator for a small programming language. Very trimmed down, it looks like this:

module Incr = Incremental.Make ()

module Ast = struct
  type t =
    | Input of int Incr.t
    | Incr of t
    | Decr of t
    | Condition of {
        if_ : t;
        then_ : t;
        else_ : t;
      }
end

let rec eval = function
  | Ast.Input x -> x
  | Incr t -> Incr.map (eval t) ~f:(( + ) 1)
  | Decr t -> Incr.map (eval t) ~f:(( - ) 1)
  | Condition { if_; then_; else_ } ->
    Incr.if_
      (Incr.map ~f:(Int.equal 0) (eval if_))
      ~then_:(eval then_) ~else_:(eval else_)

This is nice and works, we can create expressions, evaluate them, and see them react to their input changes:

let input_var = Incr.Var.create 0
let input = Incr.Var.watch input_var

let expr =
  Ast.Condition
    {
      if_ = Incr (Incr (Decr (Decr (Input input))));
      then_ = Input (Incr.return 1);
      else_ = Input (Incr.return 42);
    }

let eval_result = Incr.observe @@ eval expr

let _print_result_on_change : () =
  Incr.Observer.on_update_exn eval_result ~f:(function
    | Initialized new_value | Changed (_, new_value) ->
      Format.printf "New result: %d\n" new_value
    | _ -> ())

let set_input input_value =
  Incr.Var.set input_var input_value;
  Incr.stabilize ()

let () = set_input 1 (* Prints “New result: 42” *)
let () = set_input 0 (* Prints “New result: 1” *)
let () = set_input 1 (* Prints “New result: 42” *)
let () = set_input 2 (* Doesn't print anything because the result didn't change *)

However, the resulting incremental graph is much deeper than what I’d like, with a big tower of Map nodes corresponding to the Incr and Decr Ast nodes, with a huge impact on memory usage (I didn’t measure for this toy language, but on the real one, up to 90% of the allocations are the incremental nodes).

image
image606×1085 30 KB

What I would like instead is to be able to compress this graph into something like

image
image606×530 17.4 KB

where the Combined node corresponds to Incr.map ~f:(fun x -> x + 1 + 1 - 1 - 1)

Is there any way to achieve that without changing the structure of the eval function?

3 Likes
2

Incremental itself doesn’t provide anything like that, but you can make a module with an incremental-like interface that would fuse as desired. Something like this (not tested nor compiled, there’s probably missing type annotations because GADTs):

module Incr = ...
module Fusing_incr = struct
  type (_, _) k = Id : ('a, 'a) k | F : ('a -> 'b) -> ('a, 'b) k
  (* distinguish the identity continuation to avoid creating spurious map nodes
     in to_incr and spurious compositions in map *)

  type 'b t = T : 'a Incr.t * ('a, 'b) k -> 'b t
  let of_incr base = T (base, Id)
  let to_incr (T (base, k)) =
    match k with Id -> base | F f -> Incr.map base ~f
  let map (T (base, k)) ~f:f2 = 
    T (base, match k with Id -> f2 | F f -> (fun x -> f2 (f x)))
  let if_ a ~then_ ~else_ =
    of_incr (Incr.if_ (to_incr a) ~then_:(to_incr then_) ~else_:(to_incr else_))
end

let rec eval = function
  | Ast.Input x -> Fusing_incr.of_incr x
  | Incr t -> Fusing_incr.map (eval t) ~f:(( + ) 1)
  | Decr t -> Fusing_incr.map (eval t) ~f:(( - ) 1)
  | Condition { if_; then_; else_ } ->
    Fusing_incr.if_
      (Fusing_incr.map ~f:(Int.equal 0) (eval if_))
      ~then_:(eval then_) ~else_:(eval else_)
let eval x = Fusing_incr.to_incr (eval x)

I’ll note that if incremental nodes are used non-linearly, i.e. plugged into outer computations multiple times:

let x : _ Incr.t = Incr.map (Incr.watch (Incr.Variable.create 1)) ~f:something_expensive in
let y = Incr.map2 x x ~f:(+) in
...

naively replacing all uses of Incr by Fusing_incr:

let x : _ Fusing_incr.t = Fusing_incr.map (Fusing_incr.of_incr (Incr.watch (Incr.Variable.create 1)) ~f:something_expensive)) in
let y = Fusing_incr.both x x in
...

would cause something_expensive to run in both inputs to the Incr.both that underlies Fusing_incr.both (assuming Fusing_incr.both is implemented as fun a b -> of_incr (Incr.both (to_incr a) (to_incr b))). Once a function runs twice as often as it should, if such functions are nested, you can get exponential slowdowns.

But, to be clear, the eval function above has no such problem. And Fusing_incr could be extended to check (dynamically) that the usage is linear. Or with a slightly different interface (one where node creation has an explicit end, or at least end-of-batch), it’d probably be possible to fuse nodes used linearly, and not fuse used non-linearly, thus providing as much fusing as possible under the constraint of not duplicating any computation.

1 Like
3

Thanks for the reply @v-gb! The wrapper looks great, I’ll go implement something like that (I already need to wrap incremental in some custom arrow because there can be some big sub-graphs of the evaluation that are pure and for which I don’t want to create incremental nodes, so it will fit nicely in).

Fusing_incr could be extended to check (dynamically) that the usage is linear.

Oh, that’s interesting. How would you do that? Something like

module Fusing_incr = struct
  type (_, _) k = Id : ('a, 'a) k | F : ('a -> 'b) -> ('a, 'b) k
  (* distinguish the identity continuation to avoid creating spurious map nodes
     in to_incr and spurious compositions in map *)

  type 'b t = T : { base: 'a Incr.t; k: ('a, 'b) k; mutable consumed : bool } -> 'b t

  let of_incr base = T { base; k = Id; consumed = false }

  let consume t =
    if t.consumed then
      failwith "value is used non-linearly"
    else
      t.consumed <- true

  (* Wrap all the functions that use a `Fusing_incr.t` to call `consume` first *)
end

or did you have something smarter in mind?

4

Yeah, something like your boolean. Or it should be possible to reject only non-linear uses of the function in the F constructor, e.g. if allowing multiple calls to to_incr was useful.