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 | class LinalgStructuredOperation(IRDLOperation, ABC):
"""
Abstract base class for structured linalg operations, allowing them to be processed
via a unified interface.
"""
inputs = var_operand_def()
"""
The operands that won't be mutated.
"""
outputs = var_operand_def(ShapedType)
"""
The operands that will be accumulated into.
These inputs may be `memref`s, which will be mutated in-place, or `tensor`s, which will be returned as results.
"""
res = var_result_def(TensorType)
"""
The updated `outputs`, empty if the inputs are memrefs.
"""
body = region_def("single_block")
"""
The body implementing the combination of scalar elements of the inputs, and
yielding the scalar elements of the outputs.
"""
@abstractmethod
def get_indexing_maps(self) -> ArrayAttr[AffineMapAttr]:
"""
Get the indexing maps corresponding to this operation's operands, in order.
"""
@abstractmethod
def get_iterator_types(self) -> ArrayAttr[IteratorTypeAttr]:
"""
Get the iterator types corresponding to this operation's loop, in order.
"""
def get_num_loops(self) -> int:
return self.get_indexing_maps().data[0].data.num_dims
def get_loops_to_shapes_map(self) -> AffineMap:
"""
Returns a map to answer the question: "given an iteration space over
the codomain, what are the subshapes of the operands involved in the
computation".
The default behavior is to just concatenate all the indexing maps.
"""
indexing_maps = tuple(attr.data for attr in self.get_indexing_maps())
result_exprs = tuple(res for map in indexing_maps for res in map.results)
dims = self.get_num_loops()
# FIXME: Support symbols.
for map in indexing_maps:
if map.num_symbols != 0:
raise NotImplementedError(
"Indexing maps with symbols not supported for now."
)
syms = 0
return AffineMap(dims, syms, result_exprs)
def get_shapes_to_loops_map(self) -> AffineMap:
"""
Returns a map to answer the question: "Given a list of operand ranges,
what is the subportion of the iteration space involved in the
computation". This is the inverse problem of `get_loops_to_shapes_map`.
Return the empty AffineMap when such an AffineMap cannot be
constructed. The default behavior is based on a very simple inference
procedure that only works with permutation affine maps. A more advanced
Tensor-Comprehension like inference is possible but has proven to be
ambiguous in unfavorable case. A safer and more robust alternative is
to allow each op to define its own AffineMap.
"""
loops_to_shapes = self.get_loops_to_shapes_map()
inverse = loops_to_shapes.inverse_permutation()
if not inverse:
raise NotImplementedError(
"Non-invertible maps need dynamic shapes, which are not implemented."
)
return inverse
def get_loop_bound_sources(
self,
) -> tuple[LoopBoundSource, ...]:
"""
Return where each loop upper bound comes from.
Each entry identifies the shaped operand, the dimension index, and the size value.
"""
shapes_to_loops = self.get_shapes_to_loops_map()
needed_positions = tuple(
expr.position
for expr in shapes_to_loops.results
if isinstance(expr, AffineDimExpr)
)
assert len(shapes_to_loops.results) == len(needed_positions)
flat_shape_dims = tuple(
LoopBoundSource(operand, dim_index, dim_size)
for operand in self.operands
if isa(operand, SSAValue[ShapedType])
for dim_index, dim_size in enumerate(operand.type.get_shape())
)
return tuple(flat_shape_dims[position] for position in needed_positions)
def get_static_shapes(self) -> list[int]:
return [
dim
for operand in self.operands
if isinstance(operand.type, ShapedType)
for dim in operand.type.get_shape()
]
def get_static_loop_ranges(self) -> tuple[int, ...]:
shapes_to_loops = self.get_shapes_to_loops_map()
static_shapes = self.get_static_shapes()
return shapes_to_loops.eval(static_shapes, [])
def verify_(self) -> None:
# Operands are all tensors or all memrefs. Each kind is written back a
# different way, and MLIR requires one or the other too.
shaped_types = [
operand.type
for operand in self.operands
if isinstance(operand.type, ShapedType)
]
if any(isinstance(t, TensorType) for t in shaped_types) and any(
isinstance(t, MemRefType) for t in shaped_types
):
raise VerifyException("expected to have pure tensor or buffer semantics")
indexing_maps = tuple(attr.data for attr in self.get_indexing_maps())
if len(indexing_maps) != len(self.operands):
raise VerifyException(
f"expected the number of indexing_map ({len(indexing_maps)}) to be "
f"equal to the number of input/output operands ({len(self.operands)})"
)
# Every indexing map is over the same loops, one per iterator type, and
# has one result per dimension of its operand, which says how the loops
# reach that dimension. A scalar has none.
num_loops = len(self.get_iterator_types())
for index, (operand, indexing_map) in enumerate(
zip(self.operands, indexing_maps, strict=True)
):
if indexing_map.num_dims != num_loops:
raise VerifyException(
f"expected indexing_map #{index} to have {num_loops} dim(s) to "
f"match the number of loops"
)
if indexing_map.num_symbols:
raise VerifyException(f"unexpected symbols in indexing_map #{index}")
rank = (
len(operand.type.get_shape())
if isinstance(operand.type, ShapedType)
else 0
)
if len(indexing_map.results) != rank:
raise VerifyException(
f"expected operand #{index} of rank {rank} to match the "
f"result count of indexing_map #{index}, "
f"{len(indexing_map.results)}"
)
# The body takes one argument per operand, of that operand's element
# type, or of the operand's own type for a scalar.
block_args = self.body.block.args
if len(block_args) != len(self.operands):
raise VerifyException(
"expected as many non-induction variable region arguments as the "
f"number of input/output operands, {len(self.operands)}, but got "
f"{len(block_args)}"
)
for index, (operand, arg) in enumerate(
zip(self.operands, block_args, strict=True)
):
operand_type = operand.type
element_type: Attribute = (
operand_type.get_element_type()
if isa(operand_type, MemRefType | TensorType)
else operand_type
)
if arg.type != element_type:
raise VerifyException(
f"expected type of bb argument #{index} ({arg.type}) to match "
f"element or self type of the corresponding operand "
f"({element_type})"
)
# With nothing to index there is nothing to reach into.
if not indexing_maps:
return
# The loop ranges are read back off the operand shapes, which needs each
# loop to appear on its own as some result of some map.
loops_to_shapes = self.get_loops_to_shapes_map()
shapes_to_loops = loops_to_shapes.inverse_permutation()
if shapes_to_loops is None:
raise VerifyException(
f"invalid indexing maps are non-invertible: ({loops_to_shapes})"
)
# A shape only known at runtime gives a loop range only known then, and
# nothing can be checked against it here.
static_shapes = self.get_static_shapes()
if any(dim < 0 for dim in static_shapes):
return
end_ranges = tuple(
bound - 1 for bound in shapes_to_loops.eval(static_shapes, [])
)
start_ranges = (0,) * len(end_ranges)
# Each operand is as large as the loops reach into it. A result that is
# one loop dimension reaches every index up to that loop's range, so the
# operand dimension is exactly that. One that is an expression over the
# loops, `d0 + d1` or `d0 * 2`, reaches indices that may not fill the
# dimension, so it only has to fit.
for index, (operand, indexing_map) in enumerate(
zip(self.operands, indexing_maps, strict=True)
):
if not isinstance(operand.type, ShapedType):
continue
shape = operand.type.get_shape()
starts = indexing_map.eval(start_ranges, [])
ends = indexing_map.eval(end_ranges, [])
for dim, (start, end, size) in enumerate(
zip(starts, ends, shape, strict=True)
):
if size == 0:
continue
if min(start, end) < 0:
raise VerifyException(
f"unexpected result less than 0 at expression #{dim} in "
f"{indexing_map}"
)
inferred = max(start, end) + 1
if isinstance(indexing_map.results[dim], AffineDimExpr):
if inferred != size:
raise VerifyException(
f"inferred input/output operand #{index} has shape's "
f"dimension #{dim} to be {inferred}, but found {size}"
)
elif inferred > size:
raise VerifyException(
f"inferred input/output operand #{index} has shape's "
f"dimension #{dim} to be greater than or equal to "
f"{inferred}, but found {size}"
)
|