ninetoothed 0.10.0__py3-none-any.whl → 0.11.1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- ninetoothed/jit.py +31 -3
- ninetoothed/symbol.py +7 -0
- ninetoothed/tensor.py +79 -0
- {ninetoothed-0.10.0.dist-info → ninetoothed-0.11.1.dist-info}/METADATA +1 -1
- ninetoothed-0.11.1.dist-info/RECORD +11 -0
- ninetoothed-0.10.0.dist-info/RECORD +0 -11
- {ninetoothed-0.10.0.dist-info → ninetoothed-0.11.1.dist-info}/WHEEL +0 -0
- {ninetoothed-0.10.0.dist-info → ninetoothed-0.11.1.dist-info}/licenses/LICENSE +0 -0
ninetoothed/jit.py
CHANGED
@@ -20,6 +20,13 @@ from ninetoothed.torchifier import Torchifier
|
|
20
20
|
|
21
21
|
|
22
22
|
def make(arrangement, application, tensors):
|
23
|
+
"""Integrate the arrangement and the application of the tensors.
|
24
|
+
|
25
|
+
:param arrangement: The arrangement of the tensors.
|
26
|
+
:param application: The application of the tensors.
|
27
|
+
:param tensors: The tensors.
|
28
|
+
:return: A handle to the compute kernel.
|
29
|
+
"""
|
23
30
|
params = inspect.signature(application).parameters
|
24
31
|
types = arrangement(*tensors)
|
25
32
|
annotations = {param: type for param, type in zip(params, types)}
|
@@ -28,14 +35,26 @@ def make(arrangement, application, tensors):
|
|
28
35
|
return jit(application)
|
29
36
|
|
30
37
|
|
31
|
-
def jit(
|
38
|
+
def jit(func=None, *, _prettify=False):
|
39
|
+
"""A decorator for generating compute kernels.
|
40
|
+
|
41
|
+
:param func: The function to be compiled.
|
42
|
+
:param _prettify: Whether to prettify the generated code.
|
43
|
+
:return: A handle to the compute kernel.
|
44
|
+
|
45
|
+
.. note::
|
46
|
+
|
47
|
+
The ``_prettify`` parameter is experimental, which might break
|
48
|
+
the generated code.
|
49
|
+
"""
|
50
|
+
|
32
51
|
def wrapper(func):
|
33
52
|
return JIT(func, _prettify=_prettify)()
|
34
53
|
|
35
|
-
if
|
54
|
+
if func is None:
|
36
55
|
return wrapper
|
37
56
|
|
38
|
-
return wrapper(
|
57
|
+
return wrapper(func)
|
39
58
|
|
40
59
|
|
41
60
|
class JIT:
|
@@ -472,6 +491,15 @@ class CodeGenerator(ast.NodeTransformer):
|
|
472
491
|
for target_dim in range(tensor.target.ndim)
|
473
492
|
if offsets[source_dim][target_dim] != 0
|
474
493
|
),
|
494
|
+
) & functools.reduce(
|
495
|
+
lambda x, y: x & y,
|
496
|
+
(
|
497
|
+
indices[dim - tensor.innermost().target.ndim][
|
498
|
+
type(self)._generate_slices(tensor, target_dim)
|
499
|
+
]
|
500
|
+
< tensor.innermost().target.shape[dim]
|
501
|
+
for dim, target_dim in enumerate(tensor.innermost().target_dims)
|
502
|
+
),
|
475
503
|
)
|
476
504
|
|
477
505
|
return pointers, mask
|
ninetoothed/symbol.py
CHANGED
@@ -7,6 +7,13 @@ import ninetoothed.naming as naming
|
|
7
7
|
|
8
8
|
|
9
9
|
class Symbol:
|
10
|
+
"""A class uesed to represent a symbol.
|
11
|
+
|
12
|
+
:param expr: The expression used to construct the symbol.
|
13
|
+
:param constexpr: Whether the symbol is a constexpr.
|
14
|
+
:param mata: Whether the symbol is a meta.
|
15
|
+
"""
|
16
|
+
|
10
17
|
def __init__(self, expr, constexpr=None, meta=None):
|
11
18
|
if isinstance(expr, type(self)):
|
12
19
|
self._node = expr._node
|
ninetoothed/tensor.py
CHANGED
@@ -8,6 +8,21 @@ from ninetoothed.symbol import Symbol
|
|
8
8
|
|
9
9
|
|
10
10
|
class Tensor:
|
11
|
+
"""A class uesed to represent a symbolic tensor.
|
12
|
+
|
13
|
+
:param ndim: The number of dimensions of the tensor.
|
14
|
+
:param shape: The shape of the tensor.
|
15
|
+
:param dtype: The element type of the tensor.
|
16
|
+
:param strides: The strides of the tensor.
|
17
|
+
:param other: The values for out-of-bounds positions.
|
18
|
+
:param constexpr_shape: Whether the sizes are constexpr.
|
19
|
+
:param name: The name of the tensor.
|
20
|
+
:param source: For internal use only.
|
21
|
+
:param source_dims: For internal use only.
|
22
|
+
:param target: For internal use only.
|
23
|
+
:param target_dims: For internal use only.
|
24
|
+
"""
|
25
|
+
|
11
26
|
num_instances = 0
|
12
27
|
|
13
28
|
def __init__(
|
@@ -70,6 +85,14 @@ class Tensor:
|
|
70
85
|
type(self).num_instances += 1
|
71
86
|
|
72
87
|
def tile(self, tile_shape, strides=None, dilation=None):
|
88
|
+
"""Tiles the tensor into a hierarchical tensor.
|
89
|
+
|
90
|
+
:param tile_shape: The shape of a tile.
|
91
|
+
:param strides: The interval at which each tile is generated.
|
92
|
+
:param dilation: The spacing between tiles.
|
93
|
+
:return: A hierarchical tensor.
|
94
|
+
"""
|
95
|
+
|
73
96
|
if strides is None:
|
74
97
|
strides = [-1 for _ in tile_shape]
|
75
98
|
|
@@ -119,6 +142,12 @@ class Tensor:
|
|
119
142
|
)
|
120
143
|
|
121
144
|
def expand(self, shape):
|
145
|
+
"""Expands the specified singleton dimensions of the tensor.
|
146
|
+
|
147
|
+
:param shape: The expanded shape.
|
148
|
+
:return: The expanded tensor.
|
149
|
+
"""
|
150
|
+
|
122
151
|
# TODO: Add error handling.
|
123
152
|
return type(self)(
|
124
153
|
shape=[
|
@@ -132,9 +161,16 @@ class Tensor:
|
|
132
161
|
],
|
133
162
|
source=self.source,
|
134
163
|
source_dims=self.source_dims,
|
164
|
+
target_dims=self.target_dims,
|
135
165
|
)
|
136
166
|
|
137
167
|
def squeeze(self, dim):
|
168
|
+
"""Removes the specified singleton dimensions of the tensor.
|
169
|
+
|
170
|
+
:param dim: The dimension(s) to be squeezed.
|
171
|
+
:return: The squeezed tensor.
|
172
|
+
"""
|
173
|
+
|
138
174
|
if not isinstance(dim, tuple):
|
139
175
|
dim = (dim,)
|
140
176
|
|
@@ -149,9 +185,20 @@ class Tensor:
|
|
149
185
|
for i, source_dim in enumerate(self.source_dims)
|
150
186
|
if i not in dim
|
151
187
|
],
|
188
|
+
target_dims=[
|
189
|
+
target_dim
|
190
|
+
for i, target_dim in enumerate(self.target_dims)
|
191
|
+
if i not in dim
|
192
|
+
],
|
152
193
|
)
|
153
194
|
|
154
195
|
def permute(self, dims):
|
196
|
+
"""Permutes the dimensions of the tensor.
|
197
|
+
|
198
|
+
:param dims: The permuted ordering of the dimensions.
|
199
|
+
:return: The permuted tensor.
|
200
|
+
"""
|
201
|
+
|
155
202
|
# TODO: Add error handling.
|
156
203
|
new_shape = [None for _ in range(self.ndim)]
|
157
204
|
new_strides = [None for _ in range(self.ndim)]
|
@@ -168,9 +215,20 @@ class Tensor:
|
|
168
215
|
strides=new_strides,
|
169
216
|
source=self.source,
|
170
217
|
source_dims=new_source_dims,
|
218
|
+
target_dims=self.target_dims,
|
171
219
|
)
|
172
220
|
|
173
221
|
def flatten(self, start_dim=None, end_dim=None):
|
222
|
+
"""Flattens the specified dimensions of the tensor.
|
223
|
+
|
224
|
+
See :func:`ravel` for the differences between :func:`flatten`
|
225
|
+
and :func:`ravel`.
|
226
|
+
|
227
|
+
:param start_dim: The first dimension to flatten.
|
228
|
+
:param end_dim: The dimension after the last to flatten.
|
229
|
+
:return: The flattened tensor.
|
230
|
+
"""
|
231
|
+
|
174
232
|
# TODO: Add error handling.
|
175
233
|
if start_dim is None:
|
176
234
|
start_dim = 0
|
@@ -197,15 +255,36 @@ class Tensor:
|
|
197
255
|
leading_source_dims + (flattening_source_dims,) + trailing_source_dims
|
198
256
|
)
|
199
257
|
|
258
|
+
leading_target_dims = self.target_dims[:start_dim]
|
259
|
+
flattening_target_dims = self.target_dims[start_dim:end_dim]
|
260
|
+
trailing_target_dims = self.target_dims[end_dim:]
|
261
|
+
|
262
|
+
new_target_dims = (
|
263
|
+
leading_target_dims + (flattening_target_dims[-1],) + trailing_target_dims
|
264
|
+
)
|
265
|
+
|
200
266
|
return type(self)(
|
201
267
|
shape=new_shape,
|
202
268
|
dtype=self.dtype,
|
203
269
|
strides=new_strides,
|
204
270
|
source=self.source,
|
205
271
|
source_dims=new_source_dims,
|
272
|
+
target_dims=new_target_dims,
|
206
273
|
)
|
207
274
|
|
208
275
|
def ravel(self):
|
276
|
+
"""Flattens the hierarchy of the tensor.
|
277
|
+
|
278
|
+
:func:`ravel` differs from :func:`flatten`, which only flattens
|
279
|
+
dimensions at a single level. For example, consider a tensor
|
280
|
+
with two levels: the first level has a shape of ``(N, P, Q)``,
|
281
|
+
and the second level has a shape of ``(C, R, S)``. After
|
282
|
+
applying :func:`ravel`, the resulting tensor will have a single
|
283
|
+
flattened level with a shape of ``(N, P, Q, C, R, S)``.
|
284
|
+
|
285
|
+
:return: The raveled tensor.
|
286
|
+
"""
|
287
|
+
|
209
288
|
# TODO: Add error handling.
|
210
289
|
new_shape = []
|
211
290
|
new_strides = []
|
@@ -1,6 +1,6 @@
|
|
1
1
|
Metadata-Version: 2.4
|
2
2
|
Name: ninetoothed
|
3
|
-
Version: 0.
|
3
|
+
Version: 0.11.1
|
4
4
|
Summary: A domain-specific language based on Triton but providing higher-level abstraction.
|
5
5
|
Project-URL: Homepage, https://github.com/InfiniTensor/ninetoothed
|
6
6
|
Project-URL: Issues, https://github.com/InfiniTensor/ninetoothed/issues
|
@@ -0,0 +1,11 @@
|
|
1
|
+
ninetoothed/__init__.py,sha256=dX34sk5GA3OgWf1Jc4gJMW3UwcGcJsuG3hs3rkiqq6g,161
|
2
|
+
ninetoothed/jit.py,sha256=U3Nen5vyx69ulW7_hnRuATW86Ag9NgVgd3U02NVB20c,24430
|
3
|
+
ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
|
4
|
+
ninetoothed/naming.py,sha256=Fl0x4eDRStTpkXjJg6179ErEnY7bR5Qi0AT6RX9C3fU,951
|
5
|
+
ninetoothed/symbol.py,sha256=mN96tp-2eUxbiNfxuxtKWNSxOSdYqlcmpY2MYQ-FiEg,4993
|
6
|
+
ninetoothed/tensor.py,sha256=Pgl3t08qNDJrjbNMLltEIwMu19vnKwMUncOq4aedTjY,11983
|
7
|
+
ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
|
8
|
+
ninetoothed-0.11.1.dist-info/METADATA,sha256=DsvP5HloDDgEiHUtPEu8blUVYCBr3_VXYILUikfOE-Y,7055
|
9
|
+
ninetoothed-0.11.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
|
10
|
+
ninetoothed-0.11.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
|
11
|
+
ninetoothed-0.11.1.dist-info/RECORD,,
|
@@ -1,11 +0,0 @@
|
|
1
|
-
ninetoothed/__init__.py,sha256=dX34sk5GA3OgWf1Jc4gJMW3UwcGcJsuG3hs3rkiqq6g,161
|
2
|
-
ninetoothed/jit.py,sha256=0LeDBpSYFgPx4hatP_ZsvElsj0d9d552OKRc__L1Jvc,23460
|
3
|
-
ninetoothed/language.py,sha256=YwjlBENmmKPTnhaQ2uYbj5MwzrCAT7MLJ6VkQ6NeXJE,504
|
4
|
-
ninetoothed/naming.py,sha256=Fl0x4eDRStTpkXjJg6179ErEnY7bR5Qi0AT6RX9C3fU,951
|
5
|
-
ninetoothed/symbol.py,sha256=rZ5nXtn-U1Nw0BBRJ-kfrwmX_zCbAi76un-Z2QFaoZc,4773
|
6
|
-
ninetoothed/tensor.py,sha256=LGS9wYmPKckZjIEXsMHdclTmVhzkfrc3avOPiCQY1tU,9153
|
7
|
-
ninetoothed/torchifier.py,sha256=aDijK5UOwK2oLXDHgDo8M959rJclEI0lcfaPr7GQTXY,1012
|
8
|
-
ninetoothed-0.10.0.dist-info/METADATA,sha256=nQWkQ--AceNN3DoD-kP9_aykZYVw8LOfqf2iO63v1Ek,7055
|
9
|
-
ninetoothed-0.10.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
|
10
|
-
ninetoothed-0.10.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
|
11
|
-
ninetoothed-0.10.0.dist-info/RECORD,,
|
File without changes
|
File without changes
|