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 ..ir import *
7  from ._ods_common import get_default_loc_context as _get_default_loc_context
8
9  from typing import Any, List, Union
10except ImportError as e:
11  raise RuntimeError("Error loading imports from extension module") from e
12
13
14def _isa(obj: Any, cls: type):
15  try:
16    cls(obj)
17  except ValueError:
18    return False
19  return True
20
21
22def _is_any_of(obj: Any, classes: List[type]):
23  return any(_isa(obj, cls) for cls in classes)
24
25
26def _is_integer_like_type(type: Type):
27  return _is_any_of(type, [IntegerType, IndexType])
28
29
30def _is_float_type(type: Type):
31  return _is_any_of(type, [BF16Type, F16Type, F32Type, F64Type])
32
33
34class ConstantOp:
35  """Specialization for the constant op class."""
36
37  def __init__(self,
38               result: Type,
39               value: Union[int, float, Attribute],
40               *,
41               loc=None,
42               ip=None):
43    if isinstance(value, int):
44      super().__init__(IntegerAttr.get(result, value), loc=loc, ip=ip)
45    elif isinstance(value, float):
46      super().__init__(FloatAttr.get(result, value), loc=loc, ip=ip)
47    else:
48      super().__init__(value, loc=loc, ip=ip)
49
50  @classmethod
51  def create_index(cls, value: int, *, loc=None, ip=None):
52    """Create an index-typed constant."""
53    return cls(
54        IndexType.get(context=_get_default_loc_context(loc)),
55        value,
56        loc=loc,
57        ip=ip)
58
59  @property
60  def type(self):
61    return self.results[0].type
62
63  @property
64  def literal_value(self) -> Union[int, float]:
65    if _is_integer_like_type(self.type):
66      return IntegerAttr(self.value).value
67    elif _is_float_type(self.type):
68      return FloatAttr(self.value).value
69    else:
70      raise ValueError("only integer and float constants have literal values")
71