{{ message }}
Repository navigation
Expand file tree
/
Copy path__init__.py
More file actions
698 lines (604 loc) · 19.4 KB
/
Copy path__init__.py
File metadata and controls
698 lines (604 loc) · 19.4 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
from __future__ import annotations
__copyright__ = "Copyright (C) 2012 Andreas Kloeckner"
__license__ = """
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""
import os
from typing import TYPE_CHECKING
from pytools import strtobool
from loopy.auto_test import auto_test_vs_ref
from loopy.codegen import PreambleInfo, generate_body, generate_code, generate_code_v2
from loopy.codegen.result import CodeGenerationResult, GeneratedProgram
from loopy.diagnostic import LoopyError, LoopyWarning
from loopy.kernel import KernelState, LoopKernel
from loopy.kernel.creation import UniqueName, make_function, make_kernel
from loopy.kernel.data import (
AddressSpace,
ArrayArg,
CallMangleInfo,
ConstantArg,
GlobalArg,
ImageArg,
KernelArgument,
SubstitutionRule,
TemporaryVariable,
ValueArg,
)
from loopy.kernel.function_interface import (
CallableKernel,
InKernelCallable,
ScalarCallable,
)
from loopy.kernel.instruction import (
Assignment,
AtomicInit,
AtomicUpdate,
BarrierInstruction,
CallInstruction,
CInstruction,
HappensAfter,
InstructionBase,
LegacyStringInstructionTag,
MemoryOrdering,
MemoryScope,
MultiAssignmentBase,
NoOpInstruction,
OrderedAtomic,
UseStreamingStoreTag,
VarAtomicity,
)
from loopy.kernel.tools import (
add_and_infer_dtypes,
add_dtypes,
find_most_recent_global_barrier,
get_dot_dependency_graph,
get_global_barrier_order,
get_subkernel_to_insn_id_map,
get_subkernels,
show_dependency_graph,
)
from loopy.library.reduction import register_reduction_parser
from loopy.options import Options
from loopy.preprocess import infer_arg_descr, preprocess_kernel, preprocess_program
from loopy.schedule import (
generate_loop_schedules,
get_one_linearized_kernel,
get_one_scheduled_kernel,
linearize,
)
from loopy.statistics import (
CountGranularity,
MemAccess,
Op,
Sync,
ToCountMap,
ToCountPolynomialMap,
gather_access_footprint_bytes,
gather_access_footprints,
get_mem_access_map,
get_op_map,
get_synchronization_map,
)
from loopy.symbolic import LinearSubscript, Reduction, TaggedVariable, TypeCast
from loopy.target import ASTBuilderBase, TargetBase
from loopy.target.c import (
CFamilyTarget,
CTarget,
CWithGNULibcTarget,
ExecutableCTarget,
ExecutableCWithGNULibcTarget,
generate_header,
)
from loopy.target.cuda import CudaTarget
from loopy.target.execution import ExecutorBase
from loopy.target.ispc import ISPCTarget
from loopy.target.opencl import OpenCLTarget
from loopy.target.pyopencl import PyOpenCLTarget
from loopy.tools import Optional, clear_in_mem_caches, memoize_on_disk, t_unit_to_python
from loopy.transform.add_barrier import add_barrier
from loopy.transform.arithmetic import (
collect_common_factors_on_increment,
fold_constants,
)
from loopy.transform.batch import to_batched
from loopy.transform.buffer import buffer_array
from loopy.transform.callable import (
inline_callable_kernel,
merge,
register_callable,
rename_callable,
)
from loopy.transform.concatenate import concatenate_arrays
from loopy.transform.data import (
add_prefetch,
allocate_temporaries_for_base_storage,
change_arg_to_image,
remove_unused_arguments,
rename_argument,
set_argument_order,
set_array_axis_names,
set_array_dim_names,
set_temporary_address_space,
set_temporary_scope,
tag_array_axes,
tag_data_axes,
)
from loopy.transform.fusion import fuse_kernels
from loopy.transform.iname import (
add_inames_for_unused_hw_axes,
add_inames_to_insn,
chunk_iname,
duplicate_inames,
find_unused_axis_tag,
get_iname_duplication_options,
has_schedulable_iname_nesting,
join_inames,
make_reduction_inames_unique,
map_domain,
prioritize_loops,
remove_inames_from_insn,
remove_predicates_from_insn,
remove_unused_inames,
rename_iname,
rename_inames,
split_iname,
split_reduction_inward,
split_reduction_outward,
tag_inames,
untag_inames,
)
from loopy.transform.instruction import (
add_dependency,
add_nosync,
find_instructions,
map_instructions,
remove_instructions,
replace_instruction_ids,
set_instruction_priority,
simplify_indices,
tag_instructions,
)
from loopy.transform.loop_fusion import (
get_kennedy_unweighted_fusion_candidates,
rename_inames_in_batch,
)
from loopy.transform.pack_and_unpack_args import pack_and_unpack_args_for_call
from loopy.transform.padding import (
add_padding,
find_padding_multiple,
split_arg_axis,
split_array_axis,
split_array_dim,
)
from loopy.transform.parameter import assume, fix_parameters
from loopy.transform.precompute import precompute
from loopy.transform.privatize import (
privatize_temporaries_with_inames,
unprivatize_temporaries_with_inames,
)
from loopy.transform.realize_reduction import realize_reduction
from loopy.transform.save import save_and_reload_temporaries
from loopy.transform.subst import (
assignment_to_subst,
expand_subst,
extract_subst,
find_one_rule_matching,
find_rules_matching,
)
from loopy.translation_unit import TranslationUnit, for_each_kernel, make_program
from loopy.type_inference import infer_unknown_types
from loopy.types import LoopyType, NumpyType, ToLoopyTypeConvertible, to_loopy_type
from loopy.typing import PreambleGenerator, auto
from loopy.version import MOST_RECENT_LANGUAGE_VERSION, VERSION
if TYPE_CHECKING:
from collections.abc import Callable, Sequence
__all__ = [
"MOST_RECENT_LANGUAGE_VERSION",
"VERSION",
"ASTBuilderBase",
"AddressSpace",
"ArrayArg",
"Assignment",
"AtomicInit",
"AtomicUpdate",
"BarrierInstruction",
"CFamilyTarget",
"CInstruction",
"CTarget",
"CWithGNULibcTarget",
"CacheMode",
"CallInstruction",
"CallMangleInfo",
"CallableKernel",
"CodeGenerationResult",
"ConstantArg",
"CountGranularity",
"CudaTarget",
"ExecutableCTarget",
"ExecutableCWithGNULibcTarget",
"ExecutorBase",
"GeneratedProgram",
"GlobalArg",
"HappensAfter",
"ISPCTarget",
"ImageArg",
"InKernelCallable",
"InstructionBase",
"KernelArgument",
"KernelState",
"LegacyStringInstructionTag",
"LinearSubscript",
"LoopKernel",
"LoopyError",
"LoopyType",
"LoopyWarning",
"MemAccess",
"MemoryOrdering",
"MemoryScope",
"MultiAssignmentBase",
"NoOpInstruction",
"NumpyType",
"Op",
"OpenCLTarget",
"Optional",
"Options",
"OrderedAtomic",
"PreambleInfo",
"PyOpenCLTarget",
"Reduction",
"ScalarCallable",
"SubstitutionRule",
"Sync",
"TaggedVariable",
"TargetBase",
"TemporaryVariable",
"ToCountMap",
"ToCountPolynomialMap",
"ToLoopyTypeConvertible",
"TranslationUnit",
"TypeCast",
"UniqueName",
"UseStreamingStoreTag",
"ValueArg",
"VarAtomicity",
"add_and_infer_dtypes",
"add_barrier",
"add_dependency",
"add_dtypes",
"add_inames_for_unused_hw_axes",
"add_inames_to_insn",
"add_nosync",
"add_padding",
"add_prefetch",
"allocate_temporaries_for_base_storage",
"assignment_to_subst",
"assume",
"auto",
"auto_test_vs_ref",
"buffer_array",
"change_arg_to_image",
"chunk_iname",
"clear_in_mem_caches",
"collect_common_factors_on_increment",
"concatenate_arrays",
"duplicate_inames",
"expand_subst",
"extract_subst",
"find_instructions",
"find_most_recent_global_barrier",
"find_one_rule_matching",
"find_padding_multiple",
"find_rules_matching",
"find_unused_axis_tag",
"fix_parameters",
"fold_constants",
"for_each_kernel",
"fuse_kernels",
"gather_access_footprint_bytes",
"gather_access_footprints",
"generate_body",
"generate_code",
"generate_code_v2",
"generate_header",
"generate_loop_schedules",
"get_dot_dependency_graph",
"get_global_barrier_order",
"get_iname_duplication_options",
"get_kennedy_unweighted_fusion_candidates",
"get_mem_access_map",
"get_one_linearized_kernel",
"get_one_scheduled_kernel",
"get_op_map",
"get_subkernel_to_insn_id_map",
"get_subkernels",
"get_synchronization_map",
"has_schedulable_iname_nesting",
"infer_arg_descr",
"infer_unknown_types",
"inline_callable_kernel",
"join_inames",
"linearize",
"make_copy_kernel",
"make_einsum",
"make_function",
"make_kernel",
"make_program",
"make_reduction_inames_unique",
"map_domain",
"map_instructions",
"memoize_on_disk",
"merge",
"pack_and_unpack_args_for_call",
"precompute",
"preprocess_kernel",
"preprocess_program",
"prioritize_loops",
"privatize_temporaries_with_inames",
"realize_reduction",
"register_callable",
"register_preamble_generators",
"register_reduction_parser",
"register_symbol_manglers",
"remove_inames_from_insn",
"remove_instructions",
"remove_predicates_from_insn",
"remove_unused_arguments",
"remove_unused_inames",
"rename_argument",
"rename_callable",
"rename_iname",
"rename_inames",
"rename_inames_in_batch",
"replace_instruction_ids",
"save_and_reload_temporaries",
"set_argument_order",
"set_array_axis_names",
"set_array_dim_names",
"set_caching_enabled",
"set_instruction_priority",
"set_temporary_address_space",
"set_temporary_scope",
"show_dependency_graph",
"simplify_indices",
"split_arg_axis",
"split_array_axis",
"split_array_dim",
"split_iname",
"split_reduction_inward",
"split_reduction_outward",
"t_unit_to_python",
"tag_array_axes",
"tag_data_axes",
"tag_inames",
"tag_instructions",
"to_batched",
"to_loopy_type",
"unprivatize_temporaries_with_inames",
"untag_inames",
]
# }}}
# {{{ set_options
@for_each_kernel
def set_options(kernel, *args, **kwargs):
"""Return a new kernel with the options given as keyword arguments, or from
a string representation passed in as the first (and only) positional
argument.
See also :class:`Options`.
"""
assert isinstance(kernel, LoopKernel)
if args and kwargs:
raise TypeError("cannot pass both positional and keyword arguments")
new_opt = kernel.options.copy()
if kwargs:
from loopy.options import Options, _apply_legacy_map
kwargs = _apply_legacy_map(Options._legacy_options_map, kwargs)
for key, val in kwargs.items():
if not hasattr(new_opt, key):
raise ValueError("unknown option '%s'" % key)
setattr(new_opt, key, val)
else:
if len(args) != 1:
raise TypeError("exactly one positional argument is required if "
"no keyword args are given")
arg, = args
from loopy.options import make_options
new_opt._update(make_options(arg))
return kernel.copy(options=new_opt)
# }}}
# {{{ library registration
@for_each_kernel
def register_preamble_generators(
kernel: LoopKernel,
preamble_generators: Sequence[PreambleGenerator]) -> LoopKernel:
"""
:arg manglers: list of functions of signature ``(preamble_info)``
generating tuples ``(sortable_str_identifier, code)``,
where *preamble_info* is a :class:`PreambleInfo`.
:returns: *kernel* with *manglers* registered
"""
from loopy.tools import unpickles_equally
new_pgens = tuple(kernel.preamble_generators)
for pgen in preamble_generators:
if pgen not in new_pgens:
if not unpickles_equally(pgen):
raise LoopyError(
f"preamble generator '{pgen}' does not compare equally after "
"being unpickled and would thus disrupt loopy's caches")
new_pgens = (pgen, *new_pgens)
return kernel.copy(preamble_generators=new_pgens)
@for_each_kernel
def register_symbol_manglers(kernel, manglers):
from loopy.tools import unpickles_equally
new_manglers = kernel.symbol_manglers
for m in manglers:
if m not in new_manglers:
if not unpickles_equally(m):
raise LoopyError("mangler '%s' does not "
"compare equally after being unpickled "
"and would disrupt loopy's caches"
% m)
new_manglers = (m, *new_manglers)
return kernel.copy(symbol_manglers=new_manglers)
# }}}
# {{{ cache control
# Caching is enabled by default, but can be disabled by setting
# the environment variables LOOPY_NO_CACHE or CG_NO_CACHE to a
# 'true' value.
CACHING_ENABLED = (
not strtobool(os.environ.get("LOOPY_NO_CACHE", "false"))
and
not strtobool(os.environ.get("CG_NO_CACHE", "false")))
ABORT_ON_CACHE_MISS = strtobool(os.environ.get("LOOPY_ABORT_ON_CACHE_MISS", "false"))
def set_caching_enabled(flag):
"""Set whether :mod:`loopy` is allowed to use disk caching for its various
code generation stages.
"""
global CACHING_ENABLED
CACHING_ENABLED = flag
class CacheMode:
"""A context manager for setting whether :mod:`loopy` is allowed to use
disk caches.
"""
def __init__(self, new_flag):
self.new_flag = new_flag
def __enter__(self):
global CACHING_ENABLED
self.previous_mode = CACHING_ENABLED
CACHING_ENABLED = self.new_flag
def __exit__(self, exc_type, exc_val, exc_tb):
global CACHING_ENABLED
CACHING_ENABLED = self.previous_mode
del self.previous_mode
# }}}
# {{{ make copy kernel
def make_copy_kernel(new_dim_tags, old_dim_tags=None):
"""Returns a :class:`loopy.TranslationUnit` that changes the data layout
of a variable (called "input") to the new layout specified by
*new_dim_tags* from the one specified by *old_dim_tags*.
*old_dim_tags* defaults to an all-C layout of the same rank
as the one given by *new_dim_tags*.
"""
from loopy.kernel.array import (
SeparateArrayArrayDimTag,
VectorArrayDimTag,
parse_array_dim_tags,
)
new_dim_tags = parse_array_dim_tags(new_dim_tags, n_axes=None)
rank = len(new_dim_tags)
if old_dim_tags is None:
old_dim_tags = parse_array_dim_tags(
",".join(rank * ["c"]), n_axes=None)
elif isinstance(old_dim_tags, str):
old_dim_tags = parse_array_dim_tags(
old_dim_tags, n_axes=None)
indices = ["i%d" % i for i in range(rank)]
shape = ["n%d" % i for i in range(rank)]
command_indices = ", ".join(indices)
bounds = " and ".join(
f"0<={ind}<{shape_i}"
for ind, shape_i in zip(indices, shape, strict=True))
set_str = f"{{[{command_indices}]: {bounds} }}"
result = make_kernel(set_str,
"output[%s] = input[%s]"
% (command_indices, command_indices),
lang_version=MOST_RECENT_LANGUAGE_VERSION,
default_offset=auto)
result = tag_array_axes(result, "input", old_dim_tags)
result = tag_array_axes(result, "output", new_dim_tags)
unrolled_tags = (SeparateArrayArrayDimTag, VectorArrayDimTag)
for i in range(rank):
if (isinstance(new_dim_tags[i], unrolled_tags)
or isinstance(old_dim_tags[i], unrolled_tags)):
result = tag_inames(result, {indices[i]: "unr"})
return result
# }}}
# {{{ einsum
def make_einsum(spec, arg_names, **knl_creation_kwargs):
r"""Returns a :class:`LoopKernel` for evaluating array-based
operations using Einstein summation convention.
:param spec: a string denoting the subscripts for
summation as a comma-separated list of subscript labels.
This follows the usual :func:`numpy.einsum` convention.
Note that the explicit indicator `->` for the precise output
form is required.
:param arg_names: a sequence of string types denoting
the names of the array operands.
:param \**knl_creation_kwargs: keyword arguments for kernel creation.
See :func:`make_kernel` for a list of acceptable keyword
parameters.
.. note::
No attempt is being made to reduce the complexity
of the resulting expression. This should be dealt with
as part of a separate transformation.
"""
arg_spec, out_spec = spec.split("->")
arg_specs = arg_spec.split(",")
out_spec = out_spec.strip()
arg_specs = [arg_spec.strip() for arg_spec in arg_specs]
if len(arg_names) != len(arg_specs):
raise ValueError(
f"Number of arg names ({arg_names}) should match the number "
f"of arg specs: {arg_specs}. Length of arg names is {len(arg_names)}; "
f"expecting {len(arg_specs)} arg names."
)
out_indices = set(out_spec)
if len(out_indices) != len(out_spec):
raise ValueError(
f"Output subscripts '{out_spec}' does not contain all unique indices."
)
all_indices = {
idx
for argsp in arg_specs
for idx in argsp} | out_indices
sum_indices = all_indices - out_indices
from pymbolic import var
lhs = var("out")[tuple(var(i) for i in out_spec)]
rhs = 1
for arg_name, argsp in zip(arg_names, arg_specs, strict=True):
rhs = rhs * var(arg_name)[tuple(var(i) for i in argsp)]
if sum_indices:
rhs = Reduction("sum", tuple(var(idx) for idx in sorted(sum_indices)), rhs)
constraints = " and ".join(
"0 <= %s < N%s" % (idx, idx)
for idx in all_indices
)
knl_creation_kwargs.setdefault("name", "einsum%dto%d_kernel" % (
len(all_indices), len(out_indices)))
knl_creation_kwargs.setdefault("lang_version", MOST_RECENT_LANGUAGE_VERSION)
return make_kernel("{[%s]: %s}" % (",".join(sorted(all_indices)), constraints),
[Assignment(lhs, rhs)],
**knl_creation_kwargs)
# }}}
# {{{ default target
_DEFAULT_TARGET: Callable[[], TargetBase] | None = None
def set_default_target(target: Callable[[], TargetBase] | None):
# deliberately undocumented for now
global _DEFAULT_TARGET
_DEFAULT_TARGET = target
def _set_up_default_target():
try:
import pyopencl # ruff:ignore[unused-import]
except ImportError:
from loopy.target.opencl import OpenCLTarget
target = OpenCLTarget
else:
from loopy.target.pyopencl import PyOpenCLTarget
target = PyOpenCLTarget
set_default_target(target)
_set_up_default_target()
# }}}
# vim: foldmethod=marker
You can’t perform that action at this time.