github ↗
docs menu

Module Tolk.Op

Operations.

An operation names what a node of the intermediate representation does. One vocabulary serves every stage: tensor graphs, kernels, programs and command sequences.

Operations are totally ordered by their declaration order in t. The order is part of the contract: graph orderings and canonical operand orders are derived from it.

The arithmetic operations name their sources x, y and z, in order.

Operations

#type t =
  1. #| Special
    (*

    A hardware index, such as a workgroup or thread id. Its source bounds it and its argument names it.

    *)
  2. #| Buffer
    (*

    Storage bound to a runtime buffer, or storage private to a program: registers and workgroup memory.

    *)
  3. #| Alloc
    (*

    Storage declared without a runtime buffer. The call that owns it binds one.

    *)
  4. #| Noop
    (*

    Computes nothing.

    *)
  5. #| Param
    (*

    A parameter of a function: a buffer or a scalar variable. Its first source is its size. A view of storage has the state of storage it names as a second source (Ops.view_state).

    *)
  6. #| Refparam
    (*

    A buffer parameter passed by its address: its first source is its size, as a Param's, its second the address, a Dtype.t.Uint64.

    *)
  7. #| Call
    (*

    An invocation of the body in its first source on the remaining sources. A call that produces a value writes it into Alloc arguments.

    *)
  8. #| Program
    (*

    A program in compilation. Its sources accumulate, in order: the kernel, its Linear order, its Source and its Binary.

    *)
  9. #| Linear
    (*

    Operations in execution order.

    *)
  10. #| Source
    (*

    The rendered source text of a program.

    *)
  11. #| Binary
    (*

    A constant string of bytes: the machine code of a program, or bytes placed in a command packet or a buffer.

    *)
  12. #| Sink
    (*

    A root gathering what a graph must compute.

    *)
  13. #| After
    (*

    Its first source, ordered after the remaining sources: every consumer runs after them.

    *)
  14. #| Stack
    (*

    A vector of its sources.

    *)
  15. #| Getaddr
    (*

    The device address of a buffer or of a command sequence, for command queues.

    *)
  16. #| Index
    (*

    A pointer into its first source, offset by the others.

    *)
  17. #| Shrink
    (*

    Restricts each axis of its first source to the range between its second and third sources. In a program, a window of a buffer.

    *)
  18. #| Load
    (*

    Reads through a pointer.

    *)
  19. #| Store
    (*

    Writes its second source through its first.

    *)
  20. #| Wmma
    (*

    A matrix multiply-accumulate on matrix units.

    *)
  21. #| Cast
    (*

    Converts the value of its source to another type. Bitcast keeps the bits instead.

    *)
  22. #| Bitcast
    (*

    Reinterprets the bits as another type of the same width.

    *)
  23. #| Exp2
    (*

    2x.

    *)
  24. #| Log2
    (*

    The base-2 logarithm.

    *)
  25. #| Sin
    (*

    The sine.

    *)
  26. #| Sqrt
    (*

    The square root.

    *)
  27. #| Recip
    (*

    1 / x.

    *)
  28. #| Neg
    (*

    -x.

    *)
  29. #| Trunc
    (*

    Rounds toward zero.

    *)
  30. #| Add
    (*

    x + y.

    *)
  31. #| Mul
    (*

    x * y.

    *)
  32. #| Shl
    (*

    Shifts x left by y bits.

    *)
  33. #| Shr
    (*

    Shifts x right by y bits.

    *)
  34. #| Truncdiv
    (*

    Integer division rounding toward zero; a quotient by 0 is 0.

    *)
  35. #| Max
    (*

    The greater of x and y, computed as y where x < y and as x elsewhere. Where neither is less than the other, a NaN against any value or -0. against 0., it is either of the two: rewrites treat Max as commutative (Set.commutative) and may swap them. It is neither IEEE 754's maximum, which gives NaN, nor its maximumNumber, which gives the number; a caller that needs one tests for NaN.

    *)
  36. #| Truncmod
    (*

    The remainder of Truncdiv: it has the sign of x, and is x by 0.

    *)
  37. #| Cmplt
    (*

    x < y.

    *)
  38. #| Cmpne
    (*

    x <> y.

    *)
  39. #| Cmpeq
    (*

    x = y.

    *)
  40. #| Xor
    (*

    Bitwise exclusive or.

    *)
  41. #| Or
    (*

    Bitwise or.

    *)
  42. #| And
    (*

    Bitwise and.

    *)
  43. #| Threefry
    (*

    The Threefry-2x32 bijection of the counter x under the key y.

    *)
  44. #| Sub
    (*

    x - y.

    *)
  45. #| Fdiv
    (*

    Floating-point division.

    *)
  46. #| Pow
    (*

    xy.

    *)
  47. #| Div
    (*

    Integer division rounding toward negative infinity; a quotient by 0 is 0.

    *)
  48. #| Mod
    (*

    The remainder of Div: it has the sign of y, and is x by 0.

    *)
  49. #| Where
    (*

    y where x holds, z otherwise.

    *)
  50. #| Mulacc
    (*

    x * y + z.

    *)
  51. #| Barrier
    (*

    Synchronises the threads of a workgroup.

    *)
  52. #| Range
    (*

    A loop variable counting from 0 to its source, exclusive, or an unbounded loop when its type is void. A range of at most one trip is a conditional: a node runs inside it only when it depends on it (Ops.ranges) or reads through an After of it.

    *)
  53. #| If
    (*

    Runs the operations it guards only when its condition holds.

    *)
  54. #| End
    (*

    Closes the loops of its range sources around an effect.

    *)
  55. #| Endif
    (*

    Closes an If.

    *)
  56. #| Backedge
    (*

    Runs its body, then repeats its unbounded Range while its condition holds.

    *)
  57. #| Const
    (*

    A constant. An integer literal an operation reads holds the value that operation's integer type gives it (Ops.v).

    *)
  58. #| Ins
    (*

    An instruction named by its argument: a machine instruction, or a command of a device queue such as a wait, a barrier or a timestamp.

    *)
  59. #| Contiguous_backward
    (*

    Its source, whose gradient is made contiguous.

    *)
  60. #| Detach
    (*

    Its source, through which no gradient flows.

    *)
  61. #| Stage
    (*

    Realizes its source into a new buffer.

    *)
  62. #| Copy
    (*

    Copies its source to another device.

    *)
  63. #| Custom_function
    (*

    A function named by its argument and implemented outside the graph, such as a device submission or an external callee.

    *)
  64. #| Reshape
    (*

    Gives its source a new shape with the same element count.

    *)
  65. #| Permute
    (*

    Permutes the axes of its source.

    *)
  66. #| Expand
    (*

    Broadcasts axes of size one of its source.

    *)
  67. #| Pad
    (*

    Pads each axis of its source.

    *)
  68. #| Reduce
    (*

    Reduces axes of its source with Add, Mul or Max, in an order the compiler chooses; over no element it is the operation's identity. Since Max gives either operand for a NaN, a maximum over elements one of which is NaN is NaN, the identity or one of the elements.

    *)
#val equal : t -> t -> bool

equal o0 o1 is true iff o0 and o1 are the same operation.

#val compare : t -> t -> int

compare o0 o1 orders o0 and o1 by their declaration order in t.

#val to_int : t -> int

to_int o is the position of o in the declaration order of t, counting from 0.

#val name : t -> string

name o is the name of o in upper case, words separated by underscores: "ADD", "CONTIGUOUS_BACKWARD".

#val of_string : string -> (t, string) Stdlib.result

of_string s is the operation whose name is s, case included: of_string "ADD" is Ok Add, and of_string "add" and of_string "Ops.ADD" are errors. The error names s.

#val pp : Stdlib.Format.formatter -> t -> unit

pp formats an operation as "Ops." followed by its name: Ops.ADD.

Sets

#module Set : sig ... end

Sets of operations.