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 Scheduled : sig ... end
Sourcemodule Lowered : sig ... end