|
| 1 | +Type System |
| 2 | +=========== |
| 3 | + |
| 4 | +In hidet script, we have a type system that contains scalar types, tensor type, as well as pointer types. |
| 5 | + |
| 6 | +Scalar types |
| 7 | +------------ |
| 8 | + |
| 9 | +Hidet supports the following scalar types: |
| 10 | +- integer types: ``i8``, ``i16``, ``i32``, ``i64`` (``int8``, ``int16``, ``int32``, ``int64``) |
| 11 | +- floating point types: ``f16``, ``f32``, ``f64``, ``bf16``, ``tf32`` (``float16``, ``float32``, ``float64``, ``bfloat16``, ``tfloat32``) |
| 12 | +- boolean type: ``bool`` |
| 13 | +- complex types: ``c64``, ``c128`` (``complex64``, ``complex128``) |
| 14 | + |
| 15 | +Some types have both short names and long names. For example, ``i8`` and ``int8`` are the same type. |
| 16 | + |
| 17 | +There are also vectorized scalar types: |
| 18 | +- vectorized integer types: ``i8x4`` (``int8x4``) |
| 19 | +- vectorized float types: ``f16x2``, ``f32x4`` (``float16x2``, ``float32x4``) |
| 20 | + |
| 21 | +Tensor type |
| 22 | +----------- |
| 23 | + |
| 24 | +Hidet is designed to simplify the tensor program writing. Therefore, we have a powerful tensor type that |
| 25 | +represents a tensor with a specific element data type, shape, and memory layout. More specifically, a |
| 26 | +tensor type has the following attributes: |
| 27 | +- ``dtype``: the data type of the tensor elements, can be any scalar type. |
| 28 | +- ``shape``: a list of expressions that represents the shape of the tensor. |
| 29 | +- ``layout``: the memory layout of the tensor. |
| 30 | + |
| 31 | +The following code snippet shows how to define a tensor type: |
| 32 | + |
| 33 | +.. code-block:: |
| 34 | +
|
| 35 | + import hidet |
| 36 | + from hidet.lang import attrs, printf |
| 37 | + from hidet.lang.types import tensor, f32 |
| 38 | +
|
| 39 | + with hidet.script_module() as script_module: |
| 40 | + @hidet.script |
| 41 | + def kernel(): |
| 42 | + attrs.func_kind = 'cpu_kernel' |
| 43 | +
|
| 44 | + # by default, the layout is a row-major layout |
| 45 | + a = tensor(dtype=f32, shape=[1024, 1024]) |
| 46 | +
|
| 47 | + a[0, 0] = 0.0 |
| 48 | +
|
| 49 | + printf("a[%d, %d] = %.1f\n", 0, 0, a[0, 0]) |
| 50 | +
|
| 51 | + module = script_module.build() |
| 52 | + module() |
| 53 | +
|
| 54 | +
|
| 55 | +Tensor shape |
| 56 | +~~~~~~~~~~~~ |
| 57 | + |
| 58 | +The shape of the tensor must be determined at the compile time. Therefore, the shape of the tensor can only |
| 59 | +be defined with constant expressions. If we want to access a tensor with shape determined at runtime with |
| 60 | +variable expressions, we can use *tensor pointer* (will be discussed later). |
| 61 | + |
| 62 | + |
| 63 | +Tensor layout |
| 64 | +~~~~~~~~~~~~~ |
| 65 | + |
| 66 | +The layout of a tensor defines how to map the coordinates of a tensor element to the linear position of the |
| 67 | +element in the memory space. Generally speaking, a layout maps a :math:`n`-dimensional coordinate |
| 68 | +:math:`(c_0, c_1, \dots, c_{n-1})` to a linear index: |
| 69 | + |
| 70 | +.. math:: |
| 71 | +
|
| 72 | + index = layout(c_0, c_1, ..., c_{n-1}) |
| 73 | +
|
| 74 | +
|
| 75 | +The most commonly used layout is the row-major layout. In row-major layout, the linear index is calculated as: |
| 76 | + |
| 77 | + |
| 78 | +.. math:: |
| 79 | +
|
| 80 | + index = c_0 \times s_1 \times s_2 \times \dots \times s_{n-1} + c_1 \times s_2 \times \dots \times s_{n-1} + \dots + c_{n-2} \times s_{n-1} + c_{n-1} |
| 81 | +
|
| 82 | +where :math:`s_i` is the size of the :math:`i`-th dimension of the tensor: :math:`shape=(s_0, s_1, \dots, s_{n-1})`. |
| 83 | + |
| 84 | + |
| 85 | +Similar to the row-major layout, we can also define a column-major layout as follows: |
| 86 | + |
| 87 | +.. math:: |
| 88 | +
|
| 89 | + index = c_{n-1} \times s_{n-2} \times \dots \times s_1 \times s_0 + c_{n-2} \times \dots \times s_1 \times s_0 + \dots + c_1 \times s_0 + c_0 |
| 90 | +
|
| 91 | +The row-major layout is the default layout if we do not specify the layout of a tensor. We can also specify |
| 92 | +the layout of a tensor with the ``layout`` argument of the ``tensor`` type. For example, we can define a tensor with |
| 93 | +column-major layout as follows: |
| 94 | + |
| 95 | +.. code-block:: |
| 96 | +
|
| 97 | + from hidet.lang.layout import column_major |
| 98 | + from hidet.lang.types import tensor, f32 |
| 99 | + # ... |
| 100 | + a = tensor(dtype=f32, shape=[1024, 1024], layout=column_major(1024, 1024)) |
| 101 | + # or ignore shape if the layout is specified |
| 102 | + b = tensor(dtype=f32, layout=column_major(1024, 1024)) |
| 103 | +
|
| 104 | +
|
| 105 | +Both row-major layout and column-major layout are special cases of the strided layout. |
| 106 | +In hidet, we can define a strided layout like |
| 107 | + |
| 108 | + |
| 109 | +.. code-block:: |
| 110 | +
|
| 111 | + from hidet.lang.layout import strided_layout |
| 112 | + from hidet.lang.types import tensor, f32 |
| 113 | +
|
| 114 | + # equivalent to row-major layout |
| 115 | + a = tensor(dtype=f32, layout=strided_layout(shape=[1024, 1024], ranks=[0, 1])) |
| 116 | + # equivalent to column-major layout |
| 117 | + b = tensor(dtype=f32, layout=strided_layout(shape=[1024, 1024], ranks=[1, 0])) |
| 118 | + # the ranks define the order of the dimensions from the one that changes the slowest to the one that changes the fastest |
| 119 | + c = tensor(dtype=f32, layout=strided_layout(shape=[2, 2, 2], ranks=[0, 2, 1])) |
| 120 | + # c[coordinate] -> index |
| 121 | + # c[0, 0, 0] -> 0 |
| 122 | + # c[0, 1, 0] -> 1 |
| 123 | + # c[0, 0, 1] -> 2 |
| 124 | + # c[0, 1, 1] -> 3 |
| 125 | + # c[1, 0, 0] -> 4 |
| 126 | + # c[1, 1, 0] -> 5 |
| 127 | + # c[1, 0, 1] -> 6 |
| 128 | + # c[1, 1, 1] -> 7 |
| 129 | +
|
| 130 | +Given two layouts $f$ and $g$, we can define a new layout $h$ as the composition of $f$ and $g$ with $f$ as the outer |
| 131 | +layout and $g$ as the inner layout: |
| 132 | + |
| 133 | +.. math:: |
| 134 | +
|
| 135 | + h(\textbf{c}) = f(\textbf{c}/\textbf{s}_{g}) * n_g + g(\textbf{c} \mod \textbf{s}_{g}) |
| 136 | +
|
| 137 | +where :math:`\textbf{c}` is the coordinate of the tensor element, :math:`\textbf{s}_{g}` is the shape of the inner |
| 138 | +layout :math:`g`, and :math:`n_g` is the number of elements in the inner layout :math:`g`. The division and modulo |
| 139 | +operations are performed element-wise. The composed layout $h$ has the same number of dimensions as the outer and inner |
| 140 | +layouts, and the shape of the composed layout is the elementwise product of the shapes of the outer and inner layouts. |
| 141 | + |
| 142 | +In hidet script, we can use the *multiply* operator ``*`` to compose two layouts. For example, we can define a |
| 143 | +composed layout as follows: |
| 144 | + |
| 145 | +.. code-block:: |
| 146 | +
|
| 147 | + from hidet.lang.layout import row_major, column_major |
| 148 | +
|
| 149 | + c = row_major(2, 1) * row_major(2, 2) |
| 150 | + # c shape=[4, 2] |
| 151 | + # c[0, 0] -> 0 |
| 152 | + # c[0, 1] -> 1 |
| 153 | + # c[1, 0] -> 2 |
| 154 | + # c[1, 1] -> 3 |
| 155 | + # c[2, 0] -> 4 |
| 156 | + # c[2, 1] -> 5 |
| 157 | + # c[3, 0] -> 6 |
| 158 | + # c[3, 1] -> 7 |
| 159 | +
|
| 160 | + d = row_major(2, 1) * column_major(2, 2) |
| 161 | + # d shape=[4, 2] |
| 162 | + # d[0, 0] -> 0 |
| 163 | + # d[1, 0] -> 1 |
| 164 | + # d[0, 1] -> 2 |
| 165 | + # d[1, 1] -> 3 |
| 166 | + # d[2, 0] -> 4 |
| 167 | + # d[3, 0] -> 5 |
| 168 | + # d[2, 1] -> 6 |
| 169 | + # d[3, 1] -> 7 |
| 170 | +
|
| 171 | +We can apply the composition operation multiple times to compose multiple layouts. For example, |
| 172 | + |
| 173 | +.. code-block:: |
| 174 | +
|
| 175 | + from hidet.lang.layout import row_major, column_major |
| 176 | +
|
| 177 | + e = row_major(2, 1) * row_major(2, 2) * column_major(2, 2) # e shape=[8, 4] |
| 178 | +
|
| 179 | +The composition operation is associative, i.e., :math:`(f * g) * h = f * (g * h)`, but not commutative, i.e., |
| 180 | +it is highly likely :math:`f * g \neq g * f`. |
| 181 | + |
| 182 | + |
| 183 | +Pointer types |
| 184 | +~~~~~~~~~~~~~ |
| 185 | + |
| 186 | +In hidet, we can define a pointer type with the same semantics as the pointer type in C/C++. |
| 187 | + |
| 188 | +To construct a pointer type, we use the ``~`` operator to apply to a scalar type or pointer type: |
| 189 | + |
| 190 | +- ``~i32``: a pointer to ``i32`` type |
| 191 | +- ``~(~f16)``: a pointer to a pointer to ``f16`` type |
| 192 | + |
| 193 | + |
| 194 | +Void type |
| 195 | +~~~~~~~~~ |
| 196 | + |
| 197 | +The ``void`` type can be used as the return type of a function, or used to define a ``void`` pointer type |
| 198 | +(i.e., ``~void``). |
0 commit comments