1#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
2#  See https://llvm.org/LICENSE.txt for license information.
3#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
4"""Represents configured ops as emitted for code generation.
5
6Classes in this module generally are directly serializable to YAML for use
7by the code generator.
8
9TODO: These should just be dumb containers or serialization code but they
10currently encode too many details of how the language is interpreted. Move this
11to helpers on the comprehension objects themselves.
12"""
13
14from typing import Dict, Optional
15
16from ..... import ir as _ir
17from .comprehension import *
18from .yaml_helper import *
19
20__all__ = ["LinalgStructuredOpConfig", "LinalgOpConfig", "OperandDefConfig"]
21
22
23def _serialize_affine_map(affine_map: _ir.AffineMap) -> str:
24  with affine_map.context:
25    # Affine map printing/parsing is via an AffineMap attr.
26    attr = _ir.AffineMapAttr.get(affine_map)
27    return str(attr)
28
29
30class TensorUseConfig:
31  """Wrapper around a TensorUse with additional context-bound state."""
32
33  def __init__(self, tensor_use: TensorUse, indexing_map: _ir.AffineMap):
34    self.tensor_use = tensor_use
35    self.indexing_map = indexing_map
36
37  def __repr__(self):
38    return f"Use({self.tensor_use}, indexing_map={self.indexing_map})"
39
40
41class OperandDefConfig(YAMLObject):
42  """Wrapper containing an operand definition with additional state."""
43  yaml_tag = "!LinalgOperandDefConfig"
44
45  def __init__(self,
46               operand_def: OperandDef,
47               shape_map: Optional[_ir.AffineMap] = None,
48               index_attr_map: Optional[_ir.AffineMap] = None):
49    self.operand_def = operand_def
50    self.shape_map = shape_map  # type: Optional[_ir.AffineMap]
51    self.index_attr_map = index_attr_map  # type: Optional[_ir.AffineMap]
52    self.indexing_map = None  # type: Optional[_ir.AffineMap]
53
54  @property
55  def name(self) -> str:
56    return self.operand_def.name
57
58  @property
59  def kind(self) -> OperandKind:
60    return self.operand_def.kind
61
62  @property
63  def type_var(self) -> TypeVar:
64    return self.operand_def.type_var
65
66  def to_yaml_custom_dict(self):
67    self_dict = dict(name=self.name, kind=self.operand_def.kind.name.lower())
68    if self.type_var:
69      self_dict["type_var"] = self.type_var.name
70    if self.shape_map:
71      self_dict["shape_map"] = _serialize_affine_map(self.shape_map)
72    if self.index_attr_map:
73      self_dict["index_attr_map"] = _serialize_affine_map(self.index_attr_map)
74    if self.operand_def.default_indices:
75      self_dict["default_indices"] = self.operand_def.default_indices
76    if self.operand_def.default_fn:
77      self_dict["default_fn"] = self.operand_def.default_fn
78    return self_dict
79
80  def __repr__(self):
81    return (f"OperandDefConfig({self.operand_def}, "
82            f"shape_map={self.shape_map}, "
83            f"index_attr_map={self.index_attr_map}, "
84            f"indexing_map={self.indexing_map})")
85
86
87class LinalgIndexingMapsConfig(YAMLObject):
88  """Abstracts the style of indexing maps that the op exports.
89
90  Presently only static (tied to the op name) indexing maps are supported. In
91  the future, it is expected that we will have additional variants:
92    - Dynamic based on attributes
93    - Dynamic based on operands
94  Each is expected to require a different variant of specification.
95  """
96  yaml_tag = "!LinalgIndexingMapsConfig"
97
98  def __init__(self,
99               static_indexing_maps: Optional[Sequence[_ir.AffineMap]] = None):
100    self.static_indexing_maps = static_indexing_maps
101
102  def to_yaml_custom_dict(self):
103    if self.static_indexing_maps is not None:
104      return dict(static_indexing_maps=[
105          _serialize_affine_map(m) for m in self.static_indexing_maps
106      ])
107    raise ValueError(
108        f"LinalgIndexingMapsConfig must have one type of indexing map"
109        f"(got none)")
110
111
112class LinalgStructuredOpConfig(YAMLObject):
113  """Configuration for metadata sufficient to construct a linalg named op."""
114
115  yaml_tag = "!LinalgStructuredOpConfig"
116
117  def __init__(self,
118               comprehension: Comprehension,
119               domain: Sequence[DimDef],
120               registered_operands: Sequence[OperandDef],
121               context: Optional[_ir.Context] = None):
122    self.context = context if context is not None else _ir.Context()
123    self.affine_state = AffineBuildState()
124    self.writes = list()  # type: List[Tuple[TensorUse, TensorExpression]]
125    self.operands = dict()  # type: Dict[OperandDef, OperandDefConfig]
126    self.uses = dict()  # type: Dict[TensorUse, TensorUseConfig]
127
128    # Compute the ordered set of writes and collect the tensor, capture, dims,
129    # and index uses.
130    collected_tensor_uses = set()
131    collected_scalar_uses = set()
132    collected_dim_uses = set()
133    collected_indices = set()
134    for write_use, read_use in zip(comprehension.definitions,
135                                   comprehension.values):
136      self.writes.append((write_use, read_use))
137
138    for write_use, read_use in self.writes:
139      collected_tensor_uses.add(write_use)
140      read_use.collect_tensor_uses(collected_tensor_uses)
141      read_use.collect_scalar_uses(collected_scalar_uses)
142      read_use.collect_dim_uses(collected_dim_uses)
143      write_use.collect_dim_uses(collected_dim_uses)
144      read_use.collect_indices(collected_indices)
145
146    # Set domain to the sorted list of uses if no domain annotation is given.
147    if not domain:
148      domain = sorted(collected_dim_uses, key=lambda dim: dim.dimname)
149
150    # Verify the domain dimensions match the used dimensions.
151    if (len(domain) != len(collected_dim_uses) or
152        any(dim not in collected_dim_uses for dim in domain)):
153      raise ValueError(f"Expected the annotated domain dimensions {domain} to "
154                       f"match the set of dimension used by the tensor "
155                       f"comprehension {collected_dim_uses}")
156
157    # Instantiate the dimensions in the given order.
158    with self.context:
159      local_state = AffineBuildState(
160          global_state=self.affine_state, allow_new_symbols=False)
161      for dim in domain:
162        dim.build(state=local_state)
163
164    # Collect all attribute definitions.
165    collected_attr_defs = list()
166    for operand in registered_operands:
167      if operand.is_attribute():
168        collected_attr_defs.append(operand)
169
170    # Collect all tensors with manual indexing annotation.
171    collected_index_defs = list()
172    for operand in registered_operands:
173      if operand.index_dims:
174        if any(dim not in collected_dim_uses for dim in operand.index_dims):
175          raise ValueError(f"Expected all index dims {operand.index_dims} of "
176                           f"operand {operand.name} to have uses.")
177        collected_index_defs.append(operand)
178
179    # Collect the operand definitions of all tensor/scalar uses, attributes, and
180    # shape-only tensors.
181    all_operand_defs = list()
182    for use in collected_tensor_uses:
183      all_operand_defs.append(use.operand_def)
184    for use in collected_scalar_uses:
185      all_operand_defs.append(use.operand_def)
186    for definition in collected_attr_defs:
187      all_operand_defs.append(definition)
188    for definition in collected_index_defs:
189      all_operand_defs.append(definition)
190
191    # Add all operands in registration order to ensure the symbols are
192    # registered in the order they appear.
193    all_operand_defs = sorted(
194        all_operand_defs, key=lambda operand_def: operand_def.registered_index)
195    for operand_def in all_operand_defs:
196      self.add_operand(operand_def)
197
198    # Add all shape-only tensor index_dim annotations and all tensor uses.
199    for definition in collected_index_defs:
200      self.add_indexed_operand(definition)
201    for use in collected_tensor_uses:
202      self.add_tensor_use(use)
203
204    # Normalize all shape and indexing maps now that full count of dims and
205    # symbols are known.
206    for cuse in self.uses.values():
207      cuse.indexing_map = self._normalize_affine_map(cuse.indexing_map)
208    for definition in collected_index_defs:
209      self.operands[definition].indexing_map = self._normalize_affine_map(
210          self.operands[definition].indexing_map)
211    for operand_config in self.operands.values():
212      if operand_config.shape_map:
213        operand_config.shape_map = self._normalize_affine_map(
214            operand_config.shape_map, with_dims=False)
215      if operand_config.index_attr_map:
216        operand_config.index_attr_map = self._normalize_affine_map(
217            operand_config.index_attr_map, with_dims=False)
218
219    # Now for each write use, propagate the indexing maps from the use to the
220    # tensor, ensuring that there are not conflicts.
221    for write_use, _ in self.writes:
222      write_tensor_config = self.operands[write_use.operand_def]
223      if write_tensor_config.indexing_map:
224        raise ValueError(
225            f"Unexpected multi-write to a single tensor: {write_tensor_config}")
226      write_tensor_config.indexing_map = self.uses[write_use].indexing_map
227
228    # For each read use, propagate the indexing maps from the use to the
229    # tensor, ensuring that there are not conflicts.
230    for _, read_expr in self.writes:
231      read_uses = set()  # type: Set[TensorUse]
232      read_expr.collect_tensor_uses(read_uses)
233      for read_use in read_uses:
234        read_operand_config = self.operands[read_use.operand_def]
235        if (read_operand_config.indexing_map and
236            read_operand_config.indexing_map !=
237            self.uses[read_use].indexing_map):
238          raise ValueError(
239              f"Unexpected multi-read of a tensor with different accesses:"
240              f"{read_operand_config} vs {read_use}")
241        read_operand_config.indexing_map = self.uses[read_use].indexing_map
242
243    # Set the indexing map of all scalar uses to the empty map.
244    for operand_config in self.operands.values():
245      if operand_config.operand_def.kind == OperandKind.SCALAR:
246        operand_config.indexing_map = self._get_scalar_map()
247
248    # Check all registered tensor and scalar operands have an indexing map.
249    for operand in registered_operands:
250      if operand.is_attribute():
251        continue
252      if not (operand in self.operands and self.operands[operand].indexing_map):
253        raise ValueError(f"Failed to compute an indexing map for operand "
254                         f"{operand.name}")
255
256    # Collect reduction dims and ensure all the same.
257    all_reduction_dims = set(comprehension.all_reduction_dims)
258    if len(all_reduction_dims) != 1:
259      raise ValueError(
260          f"All writes within a generic must have the same reduction "
261          f"dims. Got: {all_reduction_dims}")
262    self.reduction_dims = next(iter(all_reduction_dims))
263
264    # Check the index dimension exists and resolve.
265    for index in collected_indices:
266      if index.dim_def.dimname not in self.affine_state.all_dims:
267        raise ValueError(
268            f"The dimension {index.dim_def.dimname} is not part of the "
269            f"iteration domain {self.affine_state.all_dims}")
270      index.resolve_dimension_name(self.affine_state)
271
272    # Generate the scalar assignments (used to build a body).
273    self.assignments = [
274        ScalarAssign(write_use.tensor_name, read_expr.to_scalar_expression())
275        for write_use, read_expr in self.writes
276    ]
277
278  @property
279  def ordered_operands(self) -> Sequence[OperandDefConfig]:
280    return sorted(
281        self.operands.values(),
282        key=lambda operand: operand.operand_def.registered_index)
283
284  @property
285  def ordered_dims(self) -> Sequence[Tuple[str, int]]:
286    """Gets the ordered list of dim bindings (symbolic name, position).
287
288    TODO: The original parser relies on parse ordering to arrive at the
289    iterator types, but that ordering is not defined on the Python side, so
290    this may be ambiguous.
291    """
292    return list(self.affine_state.all_dims.items())
293
294  @property
295  def indexing_maps(self) -> Sequence[_ir.AffineMap]:
296    return [o.indexing_map for o in self.ordered_operands if o.indexing_map]
297
298  @property
299  def iterator_types(self) -> Sequence[str]:
300
301    def get_type(symbolic_name, position):
302      for reduction_dim_expr in self.reduction_dims:
303        if reduction_dim_expr.dimname == symbolic_name:
304          return "reduction"
305      return "parallel"
306
307    return [get_type(*dim) for dim in self.ordered_dims]
308
309  def add_operand(self, operand_def: OperandDef):
310    if operand_def in self.operands:
311      return
312    if not (operand_def.is_tensor() or
313            operand_def.kind == OperandKind.INDEX_ATTR):
314      self.operands[operand_def] = OperandDefConfig(operand_def)
315      return
316    with self.context:
317      local_state = AffineBuildState(
318          global_state=self.affine_state, allow_new_dims=False)
319      exprs = []
320      for expr in operand_def.size_exprs:
321        exprs.append(expr.build(state=local_state))
322      assert local_state.local_dim_count == 0
323      affine_map = _ir.AffineMap.get(
324          dim_count=0, symbol_count=local_state.symbol_count, exprs=exprs)
325      if operand_def.kind == OperandKind.INDEX_ATTR:
326        self.operands[operand_def] = OperandDefConfig(
327            operand_def, index_attr_map=affine_map)
328      else:
329        self.operands[operand_def] = OperandDefConfig(
330            operand_def, shape_map=affine_map)
331
332  def add_indexed_operand(self, operand_def: OperandDef):
333    with self.context:
334      local_state = AffineBuildState(
335          global_state=self.affine_state, allow_new_symbols=False)
336      exprs = []
337      for expr in operand_def.index_dims:
338        exprs.append(expr.build(state=local_state))
339      self.operands[operand_def].indexing_map = _ir.AffineMap.get(
340          dim_count=local_state.dim_count,
341          symbol_count=local_state.symbol_count,
342          exprs=exprs)
343
344  def add_tensor_use(self, tensor_use: TensorUse):
345    if tensor_use in self.uses:
346      return
347    with self.context:
348      local_state = AffineBuildState(
349          global_state=self.affine_state, allow_new_symbols=False)
350      exprs = []
351      for expr in tensor_use.indices:
352        exprs.append(expr.build(state=local_state))
353      indexing_map = _ir.AffineMap.get(
354          dim_count=local_state.dim_count,
355          symbol_count=local_state.symbol_count,
356          exprs=exprs)
357
358      use_config = TensorUseConfig(tensor_use, indexing_map)
359      self.uses[tensor_use] = use_config
360
361  def _get_scalar_map(self) -> _ir.AffineMap:
362    """Create an empty affine map used to index a scalar."""
363    with self.context:
364      return _ir.AffineMap.get(
365          dim_count=self.affine_state.dim_count,
366          symbol_count=self.affine_state.symbol_count,
367          exprs=list())
368
369  def _normalize_affine_map(self,
370                            affine_map: _ir.AffineMap,
371                            with_dims: bool = True) -> _ir.AffineMap:
372    """Normalizes an indexing map to have the max known symbols and dims."""
373    with self.context:
374      return _ir.AffineMap.get(
375          dim_count=self.affine_state.dim_count if with_dims else 0,
376          symbol_count=self.affine_state.symbol_count,
377          exprs=list(affine_map.results))
378
379  def to_yaml_custom_dict(self):
380    self_dict = dict(args=self.ordered_operands)
381    # TODO: Refactor the hierarchy internally when supporting more
382    # than static (preserving this serialized form).
383    self_dict["indexing_maps"] = LinalgIndexingMapsConfig(
384        static_indexing_maps=self.indexing_maps)
385    self_dict["iterator_types"] = self.iterator_types
386    self_dict["assignments"] = self.assignments
387    return self_dict
388
389  def __repr__(self):
390    lines = [f"LinalgGenericOpConfig(reduction_dims={self.reduction_dims},"]
391    lines.append("operands=[")
392    for def_config in self.ordered_operands:
393      lines.append(f"  {repr(def_config)}")
394    lines.append("], indexing_maps=[")
395    for m in self.indexing_maps:
396      lines.append(f"  {repr(m)}")
397    lines.append(f"], iterator_types=[")
398    for t in self.iterator_types:
399      lines.append(f"  {t}")
400    lines.append("])")
401    return "\n".join(lines)
402
403
404class LinalgOpConfig(YAMLObject):
405  """Container for any supported linalg op type.
406
407  This includes the concrete type by name for ease of parsing by systems
408  that ignore tags.
409  """
410  yaml_tag = "!LinalgOpConfig"
411
412  def __init__(self,
413               metadata: OpMetadataDef,
414               *,
415               structured_op: Optional[LinalgStructuredOpConfig] = None):
416    self.metadata = metadata
417    self.structured_op = structured_op
418
419  def to_yaml_custom_dict(self):
420    self_dict = dict(metadata=self.metadata,)
421    if self.structured_op:
422      self_dict["structured_op"] = self.structured_op
423    return self_dict
424
425  @staticmethod
426  def from_linalg_op_def(
427      op_def: LinalgOpDef,
428      context: Optional[_ir.Context] = None) -> Sequence["LinalgOpConfig"]:
429    """Expands a LinalgOpDef into corresponding Linalg configured ops."""
430    # TODO: Many LinalgOpDef patterns need to expand to multiple generics.
431    assert len(op_def.comprehensions) == 1, "Only one comprehension supported"
432    return [
433        LinalgOpConfig(
434            op_def.metadata,
435            structured_op=LinalgStructuredOpConfig(
436                op_def.comprehensions[0], op_def.domain,
437                op_def.registered_operands.values(), context)),
438    ]
439
440  def __repr__(self):
441    return (f"LinalgOpConfig(metadata={self.metadata},\n"
442            f"structured_op={self.structured_op})")
443