triton-windows 3.5.0.post21__cp314-cp314-win_amd64.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.

Potentially problematic release.


This version of triton-windows might be problematic. Click here for more details.

Files changed (217) hide show
  1. triton/_C/libtriton.pyd +0 -0
  2. triton/__init__.py +82 -0
  3. triton/_filecheck.py +97 -0
  4. triton/_internal_testing.py +255 -0
  5. triton/_utils.py +126 -0
  6. triton/backends/__init__.py +47 -0
  7. triton/backends/amd/__init__.py +0 -0
  8. triton/backends/amd/compiler.py +461 -0
  9. triton/backends/amd/driver.c +283 -0
  10. triton/backends/amd/driver.py +724 -0
  11. triton/backends/amd/lib/asanrtl.bc +0 -0
  12. triton/backends/amd/lib/ockl.bc +0 -0
  13. triton/backends/amd/lib/ocml.bc +0 -0
  14. triton/backends/compiler.py +90 -0
  15. triton/backends/driver.py +66 -0
  16. triton/backends/nvidia/__init__.py +0 -0
  17. triton/backends/nvidia/bin/ptxas.exe +0 -0
  18. triton/backends/nvidia/compiler.py +533 -0
  19. triton/backends/nvidia/driver.c +517 -0
  20. triton/backends/nvidia/driver.py +799 -0
  21. triton/backends/nvidia/include/cuda.h +26280 -0
  22. triton/backends/nvidia/lib/libdevice.10.bc +0 -0
  23. triton/backends/nvidia/lib/x64/cuda.lib +0 -0
  24. triton/compiler/__init__.py +7 -0
  25. triton/compiler/code_generator.py +1614 -0
  26. triton/compiler/compiler.py +509 -0
  27. triton/compiler/errors.py +51 -0
  28. triton/compiler/make_launcher.py +0 -0
  29. triton/errors.py +5 -0
  30. triton/experimental/__init__.py +0 -0
  31. triton/experimental/gluon/__init__.py +5 -0
  32. triton/experimental/gluon/_compiler.py +0 -0
  33. triton/experimental/gluon/_runtime.py +102 -0
  34. triton/experimental/gluon/language/__init__.py +119 -0
  35. triton/experimental/gluon/language/_core.py +490 -0
  36. triton/experimental/gluon/language/_layouts.py +583 -0
  37. triton/experimental/gluon/language/_math.py +20 -0
  38. triton/experimental/gluon/language/_semantic.py +380 -0
  39. triton/experimental/gluon/language/_standard.py +80 -0
  40. triton/experimental/gluon/language/amd/__init__.py +4 -0
  41. triton/experimental/gluon/language/amd/_layouts.py +96 -0
  42. triton/experimental/gluon/language/amd/cdna3/__init__.py +100 -0
  43. triton/experimental/gluon/language/amd/cdna4/__init__.py +48 -0
  44. triton/experimental/gluon/language/amd/cdna4/async_copy.py +151 -0
  45. triton/experimental/gluon/language/extra/__init__.py +3 -0
  46. triton/experimental/gluon/language/nvidia/__init__.py +4 -0
  47. triton/experimental/gluon/language/nvidia/ampere/__init__.py +3 -0
  48. triton/experimental/gluon/language/nvidia/ampere/async_copy.py +74 -0
  49. triton/experimental/gluon/language/nvidia/ampere/mbarrier.py +80 -0
  50. triton/experimental/gluon/language/nvidia/blackwell/__init__.py +387 -0
  51. triton/experimental/gluon/language/nvidia/blackwell/tma.py +52 -0
  52. triton/experimental/gluon/language/nvidia/hopper/__init__.py +132 -0
  53. triton/experimental/gluon/language/nvidia/hopper/mbarrier.py +34 -0
  54. triton/experimental/gluon/language/nvidia/hopper/tma.py +97 -0
  55. triton/experimental/gluon/nvidia/__init__.py +4 -0
  56. triton/experimental/gluon/nvidia/blackwell.py +3 -0
  57. triton/experimental/gluon/nvidia/hopper.py +45 -0
  58. triton/knobs.py +546 -0
  59. triton/language/__init__.py +342 -0
  60. triton/language/core.py +3405 -0
  61. triton/language/extra/__init__.py +26 -0
  62. triton/language/extra/cuda/__init__.py +16 -0
  63. triton/language/extra/cuda/gdc.py +42 -0
  64. triton/language/extra/cuda/libdevice.py +1629 -0
  65. triton/language/extra/cuda/utils.py +109 -0
  66. triton/language/extra/hip/__init__.py +5 -0
  67. triton/language/extra/hip/libdevice.py +491 -0
  68. triton/language/extra/hip/utils.py +35 -0
  69. triton/language/extra/libdevice.py +790 -0
  70. triton/language/math.py +249 -0
  71. triton/language/random.py +218 -0
  72. triton/language/semantic.py +1939 -0
  73. triton/language/standard.py +534 -0
  74. triton/language/target_info.py +54 -0
  75. triton/runtime/__init__.py +23 -0
  76. triton/runtime/_allocation.py +44 -0
  77. triton/runtime/_async_compile.py +55 -0
  78. triton/runtime/autotuner.py +476 -0
  79. triton/runtime/build.py +168 -0
  80. triton/runtime/cache.py +317 -0
  81. triton/runtime/driver.py +38 -0
  82. triton/runtime/errors.py +36 -0
  83. triton/runtime/interpreter.py +1414 -0
  84. triton/runtime/jit.py +1107 -0
  85. triton/runtime/tcc/include/_mingw.h +168 -0
  86. triton/runtime/tcc/include/assert.h +62 -0
  87. triton/runtime/tcc/include/conio.h +409 -0
  88. triton/runtime/tcc/include/ctype.h +281 -0
  89. triton/runtime/tcc/include/dir.h +31 -0
  90. triton/runtime/tcc/include/direct.h +68 -0
  91. triton/runtime/tcc/include/dirent.h +135 -0
  92. triton/runtime/tcc/include/dos.h +55 -0
  93. triton/runtime/tcc/include/errno.h +75 -0
  94. triton/runtime/tcc/include/excpt.h +123 -0
  95. triton/runtime/tcc/include/fcntl.h +52 -0
  96. triton/runtime/tcc/include/fenv.h +108 -0
  97. triton/runtime/tcc/include/float.h +75 -0
  98. triton/runtime/tcc/include/inttypes.h +297 -0
  99. triton/runtime/tcc/include/io.h +418 -0
  100. triton/runtime/tcc/include/iso646.h +36 -0
  101. triton/runtime/tcc/include/limits.h +116 -0
  102. triton/runtime/tcc/include/locale.h +91 -0
  103. triton/runtime/tcc/include/malloc.h +181 -0
  104. triton/runtime/tcc/include/math.h +497 -0
  105. triton/runtime/tcc/include/mem.h +13 -0
  106. triton/runtime/tcc/include/memory.h +40 -0
  107. triton/runtime/tcc/include/process.h +176 -0
  108. triton/runtime/tcc/include/sec_api/conio_s.h +42 -0
  109. triton/runtime/tcc/include/sec_api/crtdbg_s.h +19 -0
  110. triton/runtime/tcc/include/sec_api/io_s.h +33 -0
  111. triton/runtime/tcc/include/sec_api/mbstring_s.h +52 -0
  112. triton/runtime/tcc/include/sec_api/search_s.h +25 -0
  113. triton/runtime/tcc/include/sec_api/stdio_s.h +145 -0
  114. triton/runtime/tcc/include/sec_api/stdlib_s.h +67 -0
  115. triton/runtime/tcc/include/sec_api/stralign_s.h +30 -0
  116. triton/runtime/tcc/include/sec_api/string_s.h +41 -0
  117. triton/runtime/tcc/include/sec_api/sys/timeb_s.h +34 -0
  118. triton/runtime/tcc/include/sec_api/tchar_s.h +266 -0
  119. triton/runtime/tcc/include/sec_api/time_s.h +61 -0
  120. triton/runtime/tcc/include/sec_api/wchar_s.h +128 -0
  121. triton/runtime/tcc/include/setjmp.h +160 -0
  122. triton/runtime/tcc/include/share.h +28 -0
  123. triton/runtime/tcc/include/signal.h +63 -0
  124. triton/runtime/tcc/include/stdalign.h +16 -0
  125. triton/runtime/tcc/include/stdarg.h +14 -0
  126. triton/runtime/tcc/include/stdatomic.h +171 -0
  127. triton/runtime/tcc/include/stdbool.h +11 -0
  128. triton/runtime/tcc/include/stddef.h +42 -0
  129. triton/runtime/tcc/include/stdint.h +212 -0
  130. triton/runtime/tcc/include/stdio.h +429 -0
  131. triton/runtime/tcc/include/stdlib.h +591 -0
  132. triton/runtime/tcc/include/stdnoreturn.h +7 -0
  133. triton/runtime/tcc/include/string.h +164 -0
  134. triton/runtime/tcc/include/sys/fcntl.h +13 -0
  135. triton/runtime/tcc/include/sys/file.h +14 -0
  136. triton/runtime/tcc/include/sys/locking.h +30 -0
  137. triton/runtime/tcc/include/sys/stat.h +290 -0
  138. triton/runtime/tcc/include/sys/time.h +69 -0
  139. triton/runtime/tcc/include/sys/timeb.h +133 -0
  140. triton/runtime/tcc/include/sys/types.h +123 -0
  141. triton/runtime/tcc/include/sys/unistd.h +14 -0
  142. triton/runtime/tcc/include/sys/utime.h +146 -0
  143. triton/runtime/tcc/include/tcc/tcc_libm.h +618 -0
  144. triton/runtime/tcc/include/tccdefs.h +342 -0
  145. triton/runtime/tcc/include/tcclib.h +80 -0
  146. triton/runtime/tcc/include/tchar.h +1102 -0
  147. triton/runtime/tcc/include/tgmath.h +89 -0
  148. triton/runtime/tcc/include/time.h +287 -0
  149. triton/runtime/tcc/include/uchar.h +33 -0
  150. triton/runtime/tcc/include/unistd.h +1 -0
  151. triton/runtime/tcc/include/vadefs.h +11 -0
  152. triton/runtime/tcc/include/values.h +4 -0
  153. triton/runtime/tcc/include/varargs.h +12 -0
  154. triton/runtime/tcc/include/wchar.h +873 -0
  155. triton/runtime/tcc/include/wctype.h +172 -0
  156. triton/runtime/tcc/include/winapi/basetsd.h +149 -0
  157. triton/runtime/tcc/include/winapi/basetyps.h +85 -0
  158. triton/runtime/tcc/include/winapi/guiddef.h +156 -0
  159. triton/runtime/tcc/include/winapi/poppack.h +8 -0
  160. triton/runtime/tcc/include/winapi/pshpack1.h +8 -0
  161. triton/runtime/tcc/include/winapi/pshpack2.h +8 -0
  162. triton/runtime/tcc/include/winapi/pshpack4.h +8 -0
  163. triton/runtime/tcc/include/winapi/pshpack8.h +8 -0
  164. triton/runtime/tcc/include/winapi/qos.h +72 -0
  165. triton/runtime/tcc/include/winapi/shellapi.h +59 -0
  166. triton/runtime/tcc/include/winapi/winbase.h +2958 -0
  167. triton/runtime/tcc/include/winapi/wincon.h +309 -0
  168. triton/runtime/tcc/include/winapi/windef.h +293 -0
  169. triton/runtime/tcc/include/winapi/windows.h +127 -0
  170. triton/runtime/tcc/include/winapi/winerror.h +3166 -0
  171. triton/runtime/tcc/include/winapi/wingdi.h +4080 -0
  172. triton/runtime/tcc/include/winapi/winnls.h +778 -0
  173. triton/runtime/tcc/include/winapi/winnt.h +5837 -0
  174. triton/runtime/tcc/include/winapi/winreg.h +272 -0
  175. triton/runtime/tcc/include/winapi/winsock2.h +1474 -0
  176. triton/runtime/tcc/include/winapi/winuser.h +5651 -0
  177. triton/runtime/tcc/include/winapi/winver.h +160 -0
  178. triton/runtime/tcc/include/winapi/ws2ipdef.h +21 -0
  179. triton/runtime/tcc/include/winapi/ws2tcpip.h +391 -0
  180. triton/runtime/tcc/lib/cuda.def +697 -0
  181. triton/runtime/tcc/lib/gdi32.def +337 -0
  182. triton/runtime/tcc/lib/kernel32.def +770 -0
  183. triton/runtime/tcc/lib/libtcc1.a +0 -0
  184. triton/runtime/tcc/lib/msvcrt.def +1399 -0
  185. triton/runtime/tcc/lib/python3.def +810 -0
  186. triton/runtime/tcc/lib/python310.def +1610 -0
  187. triton/runtime/tcc/lib/python311.def +1633 -0
  188. triton/runtime/tcc/lib/python312.def +1703 -0
  189. triton/runtime/tcc/lib/python313.def +1651 -0
  190. triton/runtime/tcc/lib/python313t.def +1656 -0
  191. triton/runtime/tcc/lib/python314.def +1800 -0
  192. triton/runtime/tcc/lib/python314t.def +1809 -0
  193. triton/runtime/tcc/lib/python39.def +1644 -0
  194. triton/runtime/tcc/lib/python3t.def +905 -0
  195. triton/runtime/tcc/lib/user32.def +658 -0
  196. triton/runtime/tcc/libtcc.dll +0 -0
  197. triton/runtime/tcc/tcc.exe +0 -0
  198. triton/testing.py +543 -0
  199. triton/tools/__init__.py +0 -0
  200. triton/tools/build_extern.py +365 -0
  201. triton/tools/compile.py +210 -0
  202. triton/tools/disasm.py +143 -0
  203. triton/tools/extra/cuda/compile.c +70 -0
  204. triton/tools/extra/cuda/compile.h +14 -0
  205. triton/tools/extra/hip/compile.cpp +66 -0
  206. triton/tools/extra/hip/compile.h +13 -0
  207. triton/tools/link.py +322 -0
  208. triton/tools/mxfp.py +301 -0
  209. triton/tools/ragged_tma.py +92 -0
  210. triton/tools/tensor_descriptor.py +34 -0
  211. triton/windows_utils.py +405 -0
  212. triton_windows-3.5.0.post21.dist-info/METADATA +46 -0
  213. triton_windows-3.5.0.post21.dist-info/RECORD +217 -0
  214. triton_windows-3.5.0.post21.dist-info/WHEEL +5 -0
  215. triton_windows-3.5.0.post21.dist-info/entry_points.txt +3 -0
  216. triton_windows-3.5.0.post21.dist-info/licenses/LICENSE +23 -0
  217. triton_windows-3.5.0.post21.dist-info/top_level.txt +1 -0
@@ -0,0 +1,583 @@
1
+ from dataclasses import dataclass
2
+ from typing import List, Optional
3
+ from triton.language.core import _unwrap_if_constexpr, _unwrap_shape, constexpr_type
4
+ from triton.runtime.jit import constexpr_function
5
+
6
+
7
+ def _realize_cta_layout(layout, rank):
8
+ ctas_per_cga = layout.ctas_per_cga or [1] * rank
9
+ cta_split_num = layout.cta_split_num or [1] * rank
10
+ cta_order = layout.cta_order or list(reversed(range(rank)))
11
+ object.__setattr__(layout, "ctas_per_cga", ctas_per_cga)
12
+ object.__setattr__(layout, "cta_split_num", cta_split_num)
13
+ object.__setattr__(layout, "cta_order", cta_order)
14
+
15
+
16
+ class DistributedLayout:
17
+ """
18
+ Base class for distributed memory layouts in Gluon IR.
19
+ """
20
+
21
+ @property
22
+ def type(self):
23
+ return constexpr_type(self)
24
+
25
+
26
+ @dataclass(frozen=True)
27
+ class AutoLayout(DistributedLayout):
28
+
29
+ def _to_ir(self, builder):
30
+ return builder.get_auto_layout()
31
+
32
+ def mangle(self):
33
+ return "AL"
34
+
35
+
36
+ @dataclass(frozen=True)
37
+ class BlockedLayout(DistributedLayout):
38
+ """
39
+ Represents a blocked layout, partitioning a tensor across threads, warps, and CTAs.
40
+
41
+ Args:
42
+ size_per_thread (List[int]): Number of elements per thread per dimension.
43
+ threads_per_warp (List[int]): Number of threads per warp per dimension.
44
+ warps_per_cta (List[int]): Number of warps per CTA per dimension.
45
+ order (List[int]): The ordering of dimensions for partitioning.
46
+ ctas_per_cga (Optional[List[int]]): CTAs per CGA grouping.
47
+ cta_split_num (Optional[List[int]]): Split factors for CTAs.
48
+ cta_order (Optional[List[int]]): Ordering for CTAs.
49
+ """
50
+ size_per_thread: List[int]
51
+ threads_per_warp: List[int]
52
+ warps_per_cta: List[int]
53
+ order: List[int]
54
+ ctas_per_cga: Optional[List[int]] = None
55
+ cta_split_num: Optional[List[int]] = None
56
+ cta_order: Optional[List[int]] = None
57
+
58
+ def __post_init__(self):
59
+ super().__setattr__("size_per_thread", _unwrap_if_constexpr(self.size_per_thread))
60
+ super().__setattr__("threads_per_warp", _unwrap_if_constexpr(self.threads_per_warp))
61
+ super().__setattr__("warps_per_cta", _unwrap_if_constexpr(self.warps_per_cta))
62
+ super().__setattr__("order", _unwrap_if_constexpr(self.order))
63
+ super().__setattr__("ctas_per_cga", _unwrap_if_constexpr(self.ctas_per_cga))
64
+ super().__setattr__("cta_split_num", _unwrap_if_constexpr(self.cta_split_num))
65
+ super().__setattr__("cta_order", _unwrap_if_constexpr(self.cta_order))
66
+
67
+ rank = len(self.size_per_thread)
68
+ _realize_cta_layout(self, rank)
69
+ assert len(self.threads_per_warp) == rank
70
+ assert len(self.warps_per_cta) == rank
71
+ assert len(self.order) == rank
72
+ assert len(self.ctas_per_cga) == rank
73
+ assert len(self.cta_split_num) == rank
74
+ assert len(self.cta_order) == rank
75
+
76
+ def _to_ir(self, builder):
77
+ return builder.get_blocked_layout(
78
+ self.size_per_thread,
79
+ self.threads_per_warp,
80
+ self.warps_per_cta,
81
+ self.order,
82
+ self.ctas_per_cga,
83
+ self.cta_split_num,
84
+ self.cta_order,
85
+ )
86
+
87
+ def mangle(self) -> str:
88
+
89
+ def stringify(x):
90
+ if x is None:
91
+ return ""
92
+ return "_".join(map(str, x))
93
+
94
+ size_per_thread = stringify(self.size_per_thread)
95
+ threads_per_warp = stringify(self.threads_per_warp)
96
+ warps_per_cta = stringify(self.warps_per_cta)
97
+ order = stringify(self.order)
98
+ ctas_per_cga = stringify(self.ctas_per_cga)
99
+ cta_split_num = stringify(self.cta_split_num)
100
+ cta_order = stringify(self.cta_order)
101
+ return f"B{size_per_thread}B{threads_per_warp}B{warps_per_cta}B{order}B{ctas_per_cga}B{cta_split_num}B{cta_order}B"
102
+
103
+ def __hash__(self):
104
+ return hash((
105
+ tuple(self.size_per_thread),
106
+ tuple(self.threads_per_warp),
107
+ tuple(self.warps_per_cta),
108
+ tuple(self.order),
109
+ tuple(self.ctas_per_cga) if self.ctas_per_cga else None,
110
+ tuple(self.cta_split_num) if self.cta_split_num else None,
111
+ tuple(self.cta_order) if self.cta_order else None,
112
+ ))
113
+
114
+
115
+ @dataclass(frozen=True)
116
+ class SliceLayout(DistributedLayout):
117
+ """
118
+ Represents a layout corresponding to slicing a distributed tensor along one dimension.
119
+
120
+ Args:
121
+ dim (int): The dimension index to slice.
122
+ parent (DistributedLayout): The parent layout before slicing.
123
+ """
124
+ dim: int
125
+ parent: DistributedLayout
126
+
127
+ def __post_init__(self):
128
+ super().__setattr__("dim", _unwrap_if_constexpr(self.dim))
129
+ super().__setattr__("parent", _unwrap_if_constexpr(self.parent))
130
+
131
+ def _to_ir(self, builder):
132
+ return builder.get_slice_layout(
133
+ self.dim,
134
+ self.parent._to_ir(builder),
135
+ )
136
+
137
+ def mangle(self) -> str:
138
+ return f"SL{self.dim}_{self.parent.mangle()}SL"
139
+
140
+ def __hash__(self):
141
+ return hash((self.dim, self.parent))
142
+
143
+
144
+ @dataclass(frozen=True)
145
+ class DistributedLinearLayout(DistributedLayout):
146
+ """
147
+ Represents a linear distributed layout with explicit bases at register, lane, warp, and block levels.
148
+ See: https://arxiv.org/abs/2505.23819 for reference.
149
+
150
+ Args:
151
+ reg_bases (List[List[int]]): Bases for register-level distribution.
152
+ lane_bases (List[List[int]]): Bases for lane-level distribution.
153
+ warp_bases (List[List[int]]): Bases for warp-level distribution.
154
+ block_bases (List[List[int]]): Bases for block-level distribution.
155
+ shape (List[int]): The tensor global shape.
156
+ """
157
+ reg_bases: List[List[int]]
158
+ lane_bases: List[List[int]]
159
+ warp_bases: List[List[int]]
160
+ block_bases: List[List[int]]
161
+ shape: List[int]
162
+
163
+ def __post_init__(self):
164
+ super().__setattr__("reg_bases", _unwrap_shape(self.reg_bases))
165
+ super().__setattr__("lane_bases", _unwrap_shape(self.lane_bases))
166
+ super().__setattr__("warp_bases", _unwrap_shape(self.warp_bases))
167
+ super().__setattr__("block_bases", _unwrap_shape(self.block_bases))
168
+ super().__setattr__("shape", _unwrap_shape(self.shape))
169
+
170
+ rank = len(self.shape)
171
+
172
+ for basis in self.reg_bases:
173
+ assert len(basis) == rank
174
+ for basis in self.lane_bases:
175
+ assert len(basis) == rank
176
+ for basis in self.warp_bases:
177
+ assert len(basis) == rank
178
+ for basis in self.block_bases:
179
+ assert len(basis) == rank
180
+
181
+ def _to_ir(self, builder):
182
+ return builder.get_distributed_linear_layout(self.reg_bases, self.lane_bases, self.warp_bases, self.block_bases,
183
+ self.shape)
184
+
185
+ def mangle(self):
186
+ return f"DLL{self.reg_bases}_{self.lane_bases}_{self.warp_bases}_{self.block_bases}_{self.shape}DLL"
187
+
188
+ def __hash__(self):
189
+ return hash((
190
+ tuple(map(tuple, self.reg_bases)),
191
+ tuple(map(tuple, self.lane_bases)),
192
+ tuple(map(tuple, self.warp_bases)),
193
+ tuple(map(tuple, self.block_bases)),
194
+ tuple(self.shape),
195
+ ))
196
+
197
+
198
+ @dataclass(frozen=True)
199
+ class DotOperandLayout(DistributedLayout):
200
+ """
201
+ Represents a layout for a dot operand.
202
+
203
+ Args:
204
+ operand_index (int): 0 for LHS and 1 for RHS of the dot operation.
205
+ parent (DistributedLayout): The parent layout, representing the MMA.
206
+ k_width (int): Number of elements per 32-bits.
207
+ """
208
+ operand_index: int
209
+ parent: DistributedLayout
210
+ k_width: int
211
+
212
+ def __post_init__(self):
213
+ super().__setattr__("operand_index", _unwrap_if_constexpr(self.operand_index))
214
+ super().__setattr__("parent", _unwrap_if_constexpr(self.parent))
215
+ super().__setattr__("k_width", _unwrap_if_constexpr(self.k_width))
216
+
217
+ def _to_ir(self, builder):
218
+ return builder.get_dot_operand_layout(self.operand_index, self.parent._to_ir(builder), self.k_width)
219
+
220
+ def mangle(self) -> str:
221
+ return f"DO{self.operand_index}_{self.parent.mangle()}_{self.k_width}DO"
222
+
223
+ def __hash__(self):
224
+ return hash((self.operand_index, self.parent, self.k_width))
225
+
226
+
227
+ @dataclass(frozen=True, eq=True)
228
+ class NVMMADistributedLayout(DistributedLayout):
229
+ """
230
+ Represents a layout for NVIDIA MMA (tensor core) operations.
231
+
232
+ Args:
233
+ version (List[int]): Version identifier for the MMA instruction.
234
+ warps_per_cta (List[int]): Number of warps per CTA.
235
+ instr_shape (List[int]): Instruction shape for MMA.
236
+ ctas_per_cga (Optional[List[int]]): CTAs per CGA grouping.
237
+ cta_split_num (Optional[List[int]]): Split factors for CTAs.
238
+ cta_order (Optional[List[int]]): CTA ordering.
239
+ """
240
+ version: List[int]
241
+ warps_per_cta: List[int]
242
+ instr_shape: List[int]
243
+ ctas_per_cga: Optional[List[int]] = None
244
+ cta_split_num: Optional[List[int]] = None
245
+ cta_order: Optional[List[int]] = None
246
+
247
+ def __post_init__(self):
248
+ super().__setattr__("version", _unwrap_if_constexpr(self.version))
249
+ super().__setattr__("warps_per_cta", _unwrap_if_constexpr(self.warps_per_cta))
250
+ super().__setattr__("instr_shape", _unwrap_if_constexpr(self.instr_shape))
251
+ super().__setattr__("ctas_per_cga", _unwrap_if_constexpr(self.ctas_per_cga))
252
+ super().__setattr__("cta_split_num", _unwrap_if_constexpr(self.cta_split_num))
253
+ super().__setattr__("cta_order", _unwrap_if_constexpr(self.cta_order))
254
+
255
+ rank = len(self.warps_per_cta)
256
+ _realize_cta_layout(self, rank)
257
+ assert len(self.ctas_per_cga) == rank
258
+ assert len(self.cta_split_num) == rank
259
+ assert len(self.cta_order) == rank
260
+
261
+ def _to_ir(self, builder):
262
+ return builder.get_mma_layout(self.version, self.warps_per_cta, self.ctas_per_cga, self.cta_split_num,
263
+ self.cta_order, self.instr_shape)
264
+
265
+ def mangle(self) -> str:
266
+ return f"MMA_{self.version}_{self.warps_per_cta}_{self.instr_shape}_{self.ctas_per_cga}_{self.cta_split_num}_{self.cta_order}_MMA"
267
+
268
+ def __hash__(self):
269
+ return hash((tuple(self.version), tuple(self.warps_per_cta),
270
+ tuple(self.instr_shape), tuple(self.ctas_per_cga) if self.ctas_per_cga else None,
271
+ tuple(self.cta_split_num) if self.cta_split_num else None,
272
+ tuple(self.cta_order) if self.cta_order else None))
273
+
274
+
275
+ class SharedLayout:
276
+ """
277
+ Base class for shared memory layouts in Gluon IR.
278
+ """
279
+
280
+ @property
281
+ def type(self):
282
+ return constexpr_type(self)
283
+
284
+
285
+ @constexpr_function
286
+ def _get_shape_per_cta(shape, cta_split_num):
287
+ shape_per_cta = shape
288
+ if cta_split_num is not None:
289
+ assert len(cta_split_num) == len(shape)
290
+ for dim in range(len(shape_per_cta)):
291
+ shape_per_cta[dim] /= cta_split_num[dim]
292
+ return shape_per_cta
293
+
294
+
295
+ @dataclass(frozen=True)
296
+ class NVMMASharedLayout(SharedLayout):
297
+ """
298
+ Represents a layout for shared memory suitable for NVIDIA MMA operations.
299
+
300
+ Args:
301
+ swizzle_byte_width (int): Width in bytes for swizzling.
302
+ element_bitwidth (int): Bitwidth of element type.
303
+ rank (int): Rank of the tensor.
304
+ transposed (bool): Whether the layout is transposed.
305
+ fp4_padded (bool): Whether FP4 padding is used.
306
+ ctas_per_cga (Optional[List[int]]): CTAs per CGA grouping.
307
+ cta_split_num (Optional[List[int]]): Split factors for CTAs.
308
+ cta_order (Optional[List[int]]): CTA ordering.
309
+ """
310
+ swizzle_byte_width: int
311
+ element_bitwidth: int
312
+ rank: int
313
+ transposed: bool = False
314
+ fp4_padded: bool = False
315
+ ctas_per_cga: Optional[List[int]] = None
316
+ cta_split_num: Optional[List[int]] = None
317
+ cta_order: Optional[List[int]] = None
318
+
319
+ def __post_init__(self):
320
+ super().__setattr__("swizzle_byte_width", _unwrap_if_constexpr(self.swizzle_byte_width))
321
+ super().__setattr__("element_bitwidth", _unwrap_if_constexpr(self.element_bitwidth))
322
+ super().__setattr__("rank", _unwrap_if_constexpr(self.rank))
323
+ super().__setattr__("transposed", _unwrap_if_constexpr(self.transposed))
324
+ super().__setattr__("fp4_padded", _unwrap_if_constexpr(self.fp4_padded))
325
+ super().__setattr__("ctas_per_cga", _unwrap_if_constexpr(self.ctas_per_cga))
326
+ super().__setattr__("cta_split_num", _unwrap_if_constexpr(self.cta_split_num))
327
+ super().__setattr__("cta_order", _unwrap_if_constexpr(self.cta_order))
328
+
329
+ assert self.element_bitwidth in [8, 16, 32, 64]
330
+ assert self.swizzle_byte_width in [0, 32, 64, 128]
331
+ rank = self.rank
332
+ _realize_cta_layout(self, rank)
333
+ assert len(self.ctas_per_cga) == rank
334
+ assert len(self.cta_split_num) == rank
335
+ assert len(self.cta_order) == rank
336
+
337
+ def _to_ir(self, builder):
338
+ return builder.get_nvmma_shared_layout(
339
+ self.swizzle_byte_width,
340
+ self.element_bitwidth,
341
+ self.transposed,
342
+ self.fp4_padded,
343
+ self.ctas_per_cga,
344
+ self.cta_split_num,
345
+ self.cta_order,
346
+ )
347
+
348
+ @staticmethod
349
+ @constexpr_function
350
+ def get_default_for(block_shape, dtype, transposed=False, fp4_padded=False, ctas_per_cga=None, cta_split_num=None,
351
+ cta_order=None):
352
+ """Returns an NVMMASharedLayout with default swizzling for a given shape.
353
+
354
+ This picks the largest swizzle pattern compatible with the shape, which
355
+ allows emitting the fewest TMA or MMA messages.
356
+ """
357
+ packing_factor = 2 if fp4_padded else 1
358
+ shape_per_cta = _get_shape_per_cta(block_shape, cta_split_num)
359
+ rank = len(block_shape)
360
+ if transposed:
361
+ shape_per_cta = shape_per_cta[1:] + shape_per_cta[:1]
362
+ contig_dim_size = shape_per_cta[-1] * packing_factor
363
+ contig_dim_bytes = contig_dim_size * dtype.primitive_bitwidth // 8
364
+ if contig_dim_bytes >= 128 and contig_dim_bytes % 128 == 0:
365
+ swizzle_byte_width = 128
366
+ elif contig_dim_bytes >= 64 and contig_dim_bytes % 64 == 0:
367
+ swizzle_byte_width = 64
368
+ elif contig_dim_bytes >= 32 and contig_dim_bytes % 32 == 0:
369
+ swizzle_byte_width = 32
370
+ else:
371
+ swizzle_byte_width = 0
372
+
373
+ flatten_outer_dim = 1
374
+ for size in shape_per_cta[:-1]:
375
+ flatten_outer_dim *= size
376
+ if len(block_shape) < 2 or flatten_outer_dim < 8:
377
+ swizzle_byte_width = 0
378
+
379
+ return NVMMASharedLayout(
380
+ swizzle_byte_width=swizzle_byte_width,
381
+ element_bitwidth=dtype.primitive_bitwidth,
382
+ rank=rank,
383
+ transposed=transposed,
384
+ fp4_padded=fp4_padded,
385
+ ctas_per_cga=ctas_per_cga,
386
+ cta_split_num=cta_split_num,
387
+ cta_order=cta_order,
388
+ )
389
+
390
+ def mangle(self) -> str:
391
+ return f"NVMMA_{self.swizzle_byte_width}_{self.element_bitwidth}_{self.transposed}_{self.fp4_padded}_NVMMA"
392
+
393
+ def __hash__(self):
394
+ return hash((self.swizzle_byte_width, self.element_bitwidth, self.rank, self.transposed, self.fp4_padded,
395
+ tuple(self.ctas_per_cga) if self.ctas_per_cga else None,
396
+ tuple(self.cta_split_num) if self.cta_split_num else None,
397
+ tuple(self.cta_order) if self.cta_order else None))
398
+
399
+
400
+ @dataclass(frozen=True, eq=True)
401
+ class SwizzledSharedLayout(SharedLayout):
402
+ """
403
+ Represents a generic swizzled shared memory layout.
404
+
405
+ Args:
406
+ vec (int): Vector width for swizzling.
407
+ per_phase (int): Elements per swizzle phase.
408
+ max_phase (int): Maximum number of swizzle phases.
409
+ order (List[int]): Dimension ordering for swizzling.
410
+ ctas_per_cga (Optional[List[int]]): CTAs per CGA grouping.
411
+ cta_split_num (Optional[List[int]]): Split factors for CTAs.
412
+ cta_order (Optional[List[int]]): CTA ordering.
413
+ """
414
+ vec: int
415
+ per_phase: int
416
+ max_phase: int
417
+ order: List[int]
418
+ ctas_per_cga: Optional[List[int]] = None
419
+ cta_split_num: Optional[List[int]] = None
420
+ cta_order: Optional[List[int]] = None
421
+
422
+ def __post_init__(self):
423
+ super().__setattr__("vec", _unwrap_if_constexpr(self.vec))
424
+ super().__setattr__("per_phase", _unwrap_if_constexpr(self.per_phase))
425
+ super().__setattr__("max_phase", _unwrap_if_constexpr(self.max_phase))
426
+ super().__setattr__("order", _unwrap_if_constexpr(self.order))
427
+ super().__setattr__("ctas_per_cga", _unwrap_if_constexpr(self.ctas_per_cga))
428
+ super().__setattr__("cta_split_num", _unwrap_if_constexpr(self.cta_split_num))
429
+ super().__setattr__("cta_order", _unwrap_if_constexpr(self.cta_order))
430
+
431
+ rank = len(self.order)
432
+ _realize_cta_layout(self, rank)
433
+ assert len(self.ctas_per_cga) == rank
434
+ assert len(self.cta_split_num) == rank
435
+ assert len(self.cta_order) == rank
436
+
437
+ def _to_ir(self, builder):
438
+ return builder.get_swizzled_shared_layout(
439
+ self.vec,
440
+ self.per_phase,
441
+ self.max_phase,
442
+ self.order,
443
+ self.ctas_per_cga,
444
+ self.cta_split_num,
445
+ self.cta_order,
446
+ )
447
+
448
+ def mangle(self) -> str:
449
+
450
+ def stringify(x):
451
+ if x is None:
452
+ return ""
453
+ return "_".join(map(str, x))
454
+
455
+ return f"SSS_{self.vec}_{self.per_phase}_{self.max_phase}_{stringify(self.order)}_{stringify(self.ctas_per_cga)}_{stringify(self.cta_split_num)}_{stringify(self.cta_order)}_SSS"
456
+
457
+ def __hash__(self):
458
+ return hash((self.vec, self.per_phase, self.max_phase,
459
+ tuple(self.order), tuple(self.ctas_per_cga) if self.ctas_per_cga else None,
460
+ tuple(self.cta_split_num) if self.cta_split_num else None,
461
+ tuple(self.cta_order) if self.cta_order else None))
462
+
463
+
464
+ @dataclass(frozen=True, eq=True)
465
+ class PaddedSharedLayout(SharedLayout):
466
+ """
467
+ Represents a layout for the access to shared memory. Compared to SwizzledSharedLayout,
468
+ it uses padding to avoid shared memory bank conflicts. After every interval tensor elements,
469
+ the corresponding number of padding elements are inserted.
470
+ If a position corresponds to multiple intervals, the padding amounts are summed.
471
+
472
+ In the following example of a tensor,
473
+ `eM` represents original elements in the and `pN` represents padded element.
474
+
475
+ Before padding, the shared memory looks like:
476
+ [e0, e1,
477
+ e2, e3,
478
+ e4, e5,
479
+ e6, e7,
480
+ ...]
481
+
482
+ After padding with interval-padding list [[2, 1], [4, 2]],
483
+ the shared memory will be
484
+ [e0, e1, p0,
485
+ e2, e3, p1, p2, p3,
486
+ e4, e5, p4,
487
+ e6, e7, p5, p6, p7,
488
+ ...]
489
+
490
+ Args:
491
+ interval_padding_pairs (List[int]): List of [interval, padding] pair and both interval and padding must be powers of 2.
492
+ order (List[int]): Order of logical tensor dimensions; fastest-varying first.
493
+ ctas_per_cga (Optional[List[int]]): CTAs per CGA grouping.
494
+ cta_split_num (Optional[List[int]]): Split factors for CTAs.
495
+ cta_order (Optional[List[int]]): CTA ordering.
496
+ """
497
+ interval_padding_pairs: List[List[int]]
498
+ order: List[int]
499
+ ctas_per_cga: Optional[List[int]] = None
500
+ cta_split_num: Optional[List[int]] = None
501
+ cta_order: Optional[List[int]] = None
502
+
503
+ def __post_init__(self):
504
+ super().__setattr__("interval_padding_pairs", _unwrap_shape(self.interval_padding_pairs))
505
+ super().__setattr__("order", _unwrap_if_constexpr(self.order))
506
+ super().__setattr__("ctas_per_cga", _unwrap_if_constexpr(self.ctas_per_cga))
507
+ super().__setattr__("cta_split_num", _unwrap_if_constexpr(self.cta_split_num))
508
+ super().__setattr__("cta_order", _unwrap_if_constexpr(self.cta_order))
509
+
510
+ self.verify()
511
+
512
+ def _to_ir(self, builder):
513
+ intervals, paddings = zip(*self.interval_padding_pairs)
514
+ return builder.get_padded_shared_layout(intervals, paddings, self.order, self.ctas_per_cga, self.cta_split_num,
515
+ self.cta_order)
516
+
517
+ def mangle(self) -> str:
518
+
519
+ def stringify(x):
520
+ if x is None:
521
+ return ""
522
+ return "_".join(map(str, x))
523
+
524
+ return f"PaddedShared_{stringify(self.interval_padding_pairs)}_{stringify(self.order)}_{stringify(self.ctas_per_cga)}_{stringify(self.cta_split_num)}_{stringify(self.cta_order)}_PaddedShared"
525
+
526
+ def verify(self):
527
+ pairs = self.interval_padding_pairs
528
+ assert len(pairs) > 0, "PaddedSharedLayout interval_padding_pairs must have at least one interval-padding pair"
529
+ assert all(len(pair) == 2 for pair in pairs)
530
+ intervals, paddings = zip(*pairs)
531
+
532
+ unique_intervals = list(set(intervals))
533
+ assert len(unique_intervals) == len(intervals)
534
+
535
+ is_power_of_2 = lambda n: n > 0 and n & (n - 1) == 0
536
+ assert all(is_power_of_2(n) for n in intervals), "PaddedSharedLayout interval values must all be power of two"
537
+ assert all(is_power_of_2(n) for n in paddings), "PaddedSharedLayout padding values must all be power of two"
538
+
539
+ rank = len(self.order)
540
+ assert rank > 0, "PaddedSharedLayout order must not be empty"
541
+ _realize_cta_layout(self, rank)
542
+
543
+ assert len(self.ctas_per_cga) == rank
544
+ assert len(self.cta_split_num) == rank
545
+ assert len(self.cta_order) == rank
546
+
547
+ def __hash__(self):
548
+ return hash((tuple(map(tuple, self.interval_padding_pairs)),
549
+ tuple(self.order), tuple(self.ctas_per_cga) if self.ctas_per_cga else None,
550
+ tuple(self.cta_split_num) if self.cta_split_num else None,
551
+ tuple(self.cta_order) if self.cta_order else None))
552
+
553
+
554
+ # Python impl of LinearEncodingAttr::basesPerDim
555
+ def bases_per_dim(bases, rank, skip_broadcast=True):
556
+ result = [1] * rank
557
+
558
+ if not bases:
559
+ return result
560
+
561
+ non_zero_idx = None
562
+
563
+ for basis in bases:
564
+ # Find the first non-zero index in the current basis
565
+ idx = next((i for i, v in enumerate(basis) if v != 0), None)
566
+ if idx is not None:
567
+ non_zero_idx = idx
568
+ result[idx] *= 2
569
+ elif not skip_broadcast:
570
+ # If no non-zero found and we're not skipping broadcasts, use the last found non-zero index
571
+ assert non_zero_idx is not None
572
+ result[non_zero_idx] *= 2
573
+
574
+ return result
575
+
576
+
577
+ def warps_per_cta(layout, shape):
578
+ if isinstance(layout, DistributedLinearLayout):
579
+ return bases_per_dim(layout.warp_bases, len(shape))
580
+ elif isinstance(layout, (SliceLayout, DotOperandLayout)):
581
+ return warps_per_cta(layout.parent, shape)
582
+ else:
583
+ return layout.warps_per_cta
@@ -0,0 +1,20 @@
1
+ import triton.language.math as tl_math
2
+ from ._core import builtin
3
+
4
+ umulhi = builtin(tl_math.umulhi)
5
+ exp = builtin(tl_math.exp)
6
+ exp2 = builtin(tl_math.exp2)
7
+ fma = builtin(tl_math.fma)
8
+ log = builtin(tl_math.log)
9
+ log2 = builtin(tl_math.log2)
10
+ cos = builtin(tl_math.cos)
11
+ rsqrt = builtin(tl_math.rsqrt)
12
+ sin = builtin(tl_math.sin)
13
+ sqrt = builtin(tl_math.sqrt)
14
+ sqrt_rn = builtin(tl_math.sqrt_rn)
15
+ abs = builtin(tl_math.abs)
16
+ fdiv = builtin(tl_math.fdiv)
17
+ div_rn = builtin(tl_math.div_rn)
18
+ erf = builtin(tl_math.erf)
19
+ floor = builtin(tl_math.floor)
20
+ ceil = builtin(tl_math.ceil)