package rune

  1. Overview
  2. Docs

Module Rune_jit.IrSource

Sourcemodule Dtype : sig ... end
Sourcemodule Var : sig ... end
Sourcemodule SymVar : sig ... end
Sourcemodule Special_index_kind : sig ... end
Sourcetype var_metadata = {
  1. dtype : Dtype.any;
  2. shape : int array;
  3. device : string option;
}
Sourcetype kernel_metadata = {
  1. name : string;
  2. local_dims : int;
  3. upcasted : int;
  4. dont_use_locals : bool;
}
Sourcetype custom_attr =
  1. | Attr_Int of int
  2. | Attr_Float of float
  3. | Attr_String of string
  4. | Attr_Shape of int array
Sourcetype shape_tracker = {
  1. views : view list;
  2. shape : int array;
}
Sourceand view = {
  1. shape : int array;
  2. strides : int array;
  3. offset : int;
  4. mask : (int * int) array option;
}
Sourcetype binop_kind =
  1. | Add
  2. | Mul
  3. | Sub
  4. | Div
  5. | Idiv
  6. | Fdiv
  7. | Mod
  8. | Pow
  9. | Max
  10. | Min
  11. | Cmplt
  12. | Cmpne
  13. | Xor
  14. | Or
  15. | And
  16. | Shl
  17. | Shr
Sourcetype unary_op_kind =
  1. | Neg
  2. | Log2
  3. | Exp2
  4. | Sin
  5. | Sqrt
  6. | Recip
Sourcetype ternary_op_kind =
  1. | Where
  2. | Mulacc
Sourcetype reduce_op_kind =
  1. | Reduce_Sum
  2. | Reduce_Max
  3. | Reduce_Prod
Sourcetype _ node_t =
  1. | Buffer : {
    1. dtype : 'a Dtype.t;
    2. size_in_elements : int;
    3. device : string;
    4. out_var : Var.t;
    } -> 'a node_t
  2. | Buffer_View : {
    1. buffer_var : Var.t;
    2. size : int;
    3. offset : int;
    4. dtype : 'a Dtype.t;
    5. out_var : Var.t;
    } -> 'a node_t
  3. | Placeholder : {
    1. out_var : Var.t;
    2. dtype : 'a Dtype.t;
    3. shape : int array;
    } -> 'a node_t
  4. | Const_Scalar : {
    1. value : 'a;
    2. dtype : 'a Dtype.t;
    3. out_var : Var.t;
    } -> 'a node_t
  5. | Vconst : {
    1. values : 'a array;
    2. dtype : 'a Dtype.t;
    3. out_var : Var.t;
    } -> 'a node_t
  6. | Binop : {
    1. op : binop_kind;
    2. a_var : Var.t;
    3. b_var : Var.t;
    4. out_var : Var.t;
    5. dtype : 'a Dtype.t;
    } -> 'a node_t
  7. | Unary : {
    1. op : unary_op_kind;
    2. in_var : Var.t;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  8. | Ternary : {
    1. op : ternary_op_kind;
    2. a_var : Var.t;
    3. b_var : Var.t;
    4. c_var : Var.t;
    5. out_var : Var.t;
    6. dtype : 'a Dtype.t;
    } -> 'a node_t
  9. | View : {
    1. in_var : Var.t;
    2. shape_tracker : shape_tracker;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  10. | Reshape : {
    1. in_var : Var.t;
    2. new_shape : int array;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  11. | Permute : {
    1. in_var : Var.t;
    2. axes_permutation : int array;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  12. | Expand : {
    1. in_var : Var.t;
    2. new_target_shape : int array;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  13. | Pad : {
    1. in_var : Var.t;
    2. pad_width : (int * int) array;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  14. | Shrink : {
    1. in_var : Var.t;
    2. limits : (int * int) array;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  15. | Flip : {
    1. in_var : Var.t;
    2. axes : int array;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  16. | Reduce_Axis : {
    1. in_var : Var.t;
    2. reduce_op_kind : reduce_op_kind;
    3. axes : int array;
    4. out_var : Var.t;
    5. dtype : 'a Dtype.t;
    } -> 'a node_t
  17. | Valid : {
    1. in_var : Var.t;
    2. shape_tracker : shape_tracker;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  18. | Index : {
    1. in_var : Var.t;
    2. idx_var : Var.t;
    3. valid_var : Var.t option;
    4. out_var : Var.t;
    5. dtype : 'a Dtype.t;
    } -> 'a node_t
  19. | Gep : {
    1. in_var : Var.t;
    2. indices : int array;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  20. | Vectorize : {
    1. in_vars : Var.t array;
    2. out_var : Var.t;
    3. dtype : 'a Dtype.t;
    } -> 'a node_t
  21. | Wmma : {
    1. a_var : Var.t;
    2. b_var : Var.t;
    3. c_var : Var.t;
    4. m : int;
    5. n : int;
    6. k : int;
    7. out_var : Var.t;
    8. dtype : 'a Dtype.t;
    } -> 'a node_t
  22. | Cast : {
    1. in_var : Var.t;
    2. target_dtype : Dtype.any;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  23. | Bitcast : {
    1. in_var : Var.t;
    2. target_dtype : Dtype.any;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  24. | Contiguous : {
    1. in_var : Var.t;
    2. out_var : Var.t;
    3. dtype : 'a Dtype.t;
    } -> 'a node_t
  25. | Copy : {
    1. in_var : Var.t;
    2. target_device : string;
    3. clone : bool;
    4. out_var : Var.t;
    5. dtype : 'a Dtype.t;
    } -> 'a node_t
  26. | Assign : {
    1. target_var : Var.t;
    2. updates : (Var.t * Var.t * (int * int) option) array;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  27. | Define_Var : {
    1. sym_var : SymVar.t;
    2. out_var : Var.t;
    3. dtype : 'a Dtype.t;
    } -> 'a node_t
  28. | Bind : {
    1. sym_var : Var.t;
    2. value : int;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  29. | Detach : {
    1. in_var : Var.t;
    2. out_var : Var.t;
    3. dtype : 'a Dtype.t;
    } -> 'a node_t
  30. | Contiguous_Backward : {
    1. in_var : Var.t;
    2. out_var : Var.t;
    3. dtype : 'a Dtype.t;
    } -> 'a node_t
  31. | Sink : {
    1. deps : Var.t array;
    2. dtype : 'a Dtype.t;
    } -> 'a node_t
  32. | Kernel : {
    1. ast : any_node;
    2. input_vars : Var.t array;
    3. output_vars : Var.t array;
    4. metadata : kernel_metadata;
    5. out_var : Var.t;
    6. dtype : 'a Dtype.t;
    } -> 'a node_t
  33. | Unique : {
    1. id : int;
    2. out_var : Var.t;
    3. dtype : 'a Dtype.t;
    } -> 'a node_t
  34. | Device : {
    1. device_name : string;
    2. out_var : Var.t;
    3. dtype : 'a Dtype.t;
    } -> 'a node_t
  35. | Multi : {
    1. device_vars : Var.t array;
    2. axis : int option;
    3. real_mask : bool array;
    4. out_var : Var.t;
    5. dtype : 'a Dtype.t;
    } -> 'a node_t
  36. | Fuse : {
    1. in_var : Var.t;
    2. out_var : Var.t;
    3. dtype : 'a Dtype.t;
    } -> 'a node_t
  37. | Unroll : {
    1. loop_var : Var.t;
    2. unroll_factor : int;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  38. | Contract : {
    1. in_vars : Var.t array;
    2. contraction_axes : (int * int) array;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  39. | Cat : {
    1. in_vars : Var.t array;
    2. axis : int;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  40. | Threefry : {
    1. ctr_var : Var.t;
    2. key_var : Var.t;
    3. out_var : Var.t;
    4. dtype : 'a Dtype.t;
    } -> 'a node_t
  41. | Gather : {
    1. src_var : Var.t;
    2. indices_var : Var.t;
    3. axis : int;
    4. out_var : Var.t;
    5. dtype : 'a Dtype.t;
    } -> 'a node_t
  42. | Scatter : {
    1. indices_var : Var.t;
    2. updates_var : Var.t;
    3. axis : int;
    4. shape : int array;
    5. out_var : Var.t;
    6. dtype : 'a Dtype.t;
    } -> 'a node_t
  43. | Custom : {
    1. op_name : string;
    2. in_vars : Var.t array;
    3. attributes : (string * custom_attr) list;
    4. out_var : Var.t;
    5. dtype : 'a Dtype.t;
    } -> 'a node_t
  44. | Noop : {
    1. in_var : Var.t option;
    2. out_var : Var.t;
    3. dtype : 'a Dtype.t;
    } -> 'a node_t
Sourceand any_node =
  1. | Any_Node : 'a node_t -> any_node
Sourcetype graph_t = {
  1. nodes : any_node list;
  2. vars_metadata : (Var.t, var_metadata) Hashtbl.t;
  3. input_vars : Var.t list;
  4. output_vars : Var.t list;
  5. symbolic_vars : SymVar.t list;
}
Sourceval buffer : dtype:'a Dtype.t -> size:int -> device:string -> out_var:Var.t -> 'a node_t
Sourceval unary : op:unary_op_kind -> in_var:Var.t -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval binary : op:binop_kind -> a_var:Var.t -> b_var:Var.t -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval ternary : op:ternary_op_kind -> a_var:Var.t -> b_var:Var.t -> c_var:Var.t -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval const_scalar : value:'a -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval vconst : values:'a array -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval reshape : in_var:Var.t -> new_shape:int array -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval permute : in_var:Var.t -> axes_permutation:int array -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval expand : in_var:Var.t -> new_target_shape:int array -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval pad : in_var:Var.t -> pad_width:(int * int) array -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval shrink : in_var:Var.t -> limits:(int * int) array -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval reduce_axis : in_var:Var.t -> reduce_op_kind:reduce_op_kind -> axes:int array -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval cast : in_var:Var.t -> target_dtype:Dtype.any -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval bitcast : in_var:Var.t -> target_dtype:Dtype.any -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval view : in_var:Var.t -> shape_tracker:shape_tracker -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval copy : in_var:Var.t -> target_device:string -> clone:bool -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval cat : in_vars:Var.t array -> axis:int -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval gather : src_var:Var.t -> indices_var:Var.t -> axis:int -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval scatter : indices_var:Var.t -> updates_var:Var.t -> axis:int -> shape:int array -> out_var:Var.t -> dtype:'a Dtype.t -> 'a node_t
Sourceval fresh_var : unit -> int
Sourcemodule Lowered : sig ... end
OCaml

Innovation. Community. Security.

GitHub Discord Twitter Peertube RSS
About
Changelog Releases Industrial Users Academic Users Why OCaml
Ecosystem
Packages Community Events OCaml Planet Jobs
Resources
Install OCaml Get Started Platform Tools Language Manual Standard Library API Books Exercises Papers OCaml Playground Logo
Policies
Carbon Footprint Governance Privacy Code of Conduct