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
5try:
6  from typing import Optional, Sequence, Union
7  from ..ir import *
8  from ._ods_common import get_default_loc_context
9  from .._mlir_libs._mlirDialectsLinalg import fill_builtin_region
10except ImportError as e:
11  raise RuntimeError("Error loading imports from extension module") from e
12
13from ._ods_common import get_op_result_or_value as _get_op_result_or_value
14
15def isa(cls: Type, ty: Type):
16  try:
17    cls(ty)
18    return True
19  except ValueError:
20    return False
21
22
23class InitTensorOp:
24  """Extends the linalg.init_tensor op."""
25
26  def __init__(self,
27               sizes: Union[Sequence[int], Sequence[Value]],
28               element_type: Type,
29               *,
30               loc=None,
31               ip=None):
32    """Constructs an `init_tensor` with either static or dynamic sizes."""
33    context = get_default_loc_context(loc)
34    operands = []
35    attributes = {}
36    # TODO: Refactor the InitTensorOp to take an element type attribute and
37    # then use normal result type inference, unifying the Python and C++ side
38    # with a standard mechanism (versus stashing that in builders).
39    if sizes and isinstance(sizes[0], Value):
40      # Dynamic sizes.
41      operands.extend(sizes)
42      static_size_ints = [-1] * len(sizes)
43      result_type = RankedTensorType.get(static_size_ints, element_type)
44    else:
45      # Static sizes.
46      result_type = RankedTensorType.get(sizes, element_type)
47      static_size_ints = sizes
48
49    i64_type = IntegerType.get_signless(64)
50    attributes["static_sizes"] = ArrayAttr.get(
51        [IntegerAttr.get(i64_type, s) for s in static_size_ints],
52        context=context)
53    op = self.build_generic(results=[result_type],
54                            operands=operands,
55                            attributes=attributes,
56                            loc=loc,
57                            ip=ip)
58    OpView.__init__(self, op)
59
60
61class StructuredOpMixin:
62  """All structured ops use the same mixin class."""
63
64  def __init__(self, inputs, outputs=(), results=(), loc=None, ip=None):
65    super().__init__(
66        self.build_generic(results=list(results),
67                           operands=[list(inputs), list(outputs)],
68                           loc=loc,
69                           ip=ip))
70
71
72def select_opview_mixin(parent_opview_cls):
73  # TODO: This shouldn't be a heuristic: we should have a way to annotate
74  # the OpView to note that it is a structured op.
75  if ("__init__" not in parent_opview_cls.__dict__ and
76      hasattr(parent_opview_cls, "inputs") and
77      hasattr(parent_opview_cls, "outputs") and
78      hasattr(parent_opview_cls, "result_tensors")):
79    return StructuredOpMixin
80