
# Autogenerated by mlir-tblgen; don't manually edit.

from ._ods_common import _cext as _ods_cext
from ._ods_common import (
    equally_sized_accessor as _ods_equally_sized_accessor,
    get_default_loc_context as _ods_get_default_loc_context,
    get_op_results_or_values as _get_op_results_or_values,
    segmented_accessor as _ods_segmented_accessor,
)
_ods_ir = _ods_cext.ir
_ods_cext.globals.register_traceback_file_exclusion(__file__)

import builtins
from typing import Any as _Any, Sequence as _Sequence, Union as _Union, Optional as _Optional
import sys as _sys
if _sys.version_info >= (3, 12):
  from collections.abc import Buffer as _Buffer  # pytype: disable=not-supported-yet
else:
  try:
    from typing_extensions import Buffer as _Buffer
  except ImportError:
    _Buffer = _Any


from ._transform_ops_gen import _Dialect

@_ods_cext.register_operation(_Dialect)
class AlternativesOp(_ods_ir.OpView):
  r"""
  This op represents a choice over which of its regions is to be used.
  
  When `selected_region` is provided, the semantics are that this op is to be
  substituted for by the selected region, meaning the region's results become
  the results of this op. Without a provided `selected_region`, the semantics
  are that this non-deterministic choice is yet to be resolved -- which in
  terms of the op's interpreted semantics is a failure.
  
  The `selected_region` argument is either an `IntegerAttr` or a param holding
  an `IntegerAttr`, which should provide a valid zero-based index with respect
  to the number of alternatives, i.e. regions.
  """

  OPERATION_NAME = "transform.tune.alternatives"

  _ODS_REGIONS = (0, False)

  def __init__(self, results_: _Sequence[_ods_ir.Type], name: _Union[_Any, _ods_ir.StringAttr], num_alternatives: int, *, selected_region_attr: _Optional[_Union[_Any, _ods_ir.IntegerAttr]] = None, selected_region_param: _Optional[_ods_ir.Value] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    if selected_region_param is not None: operands.append(selected_region_param)
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["name"] = (name if (
    isinstance(name, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('Builtin_StringAttr')) else
      _ods_ir.AttrBuilder.get('Builtin_StringAttr')(name, context=_ods_context))
    if selected_region_attr is not None: attributes["selected_region_attr"] = (selected_region_attr if (
        isinstance(selected_region_attr, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('APIntAttr')) else
          _ods_ir.AttrBuilder.get('APIntAttr')(selected_region_attr, context=_ods_context))
    results = []
    results.extend(results_)
    _ods_successors = None
    regions = 0 + num_alternatives
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def selected_region_param(self) -> _Optional[_ods_ir.Value]:
    return None if len(self.operation.operands) < 1 else self.operation.operands[0]

  @builtins.property
  def name(self) -> _ods_ir.StringAttr:
    return self.operation.attributes["name"]

  @name.setter
  def name(self, value: _ods_ir.StringAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["name"] = value

  @builtins.property
  def selected_region_attr(self) -> _Optional[_ods_ir.IntegerAttr]:
    if "selected_region_attr" not in self.operation.attributes:
      return None
    return self.operation.attributes["selected_region_attr"]

  @selected_region_attr.setter
  def selected_region_attr(self, value: _Optional[_ods_ir.IntegerAttr]):
    if value is not None:
      self.operation.attributes["selected_region_attr"] = value
    elif "selected_region_attr" in self.operation.attributes:
      del self.operation.attributes["selected_region_attr"]

  @selected_region_attr.deleter
  def selected_region_attr(self):
    del self.operation.attributes["selected_region_attr"]

  @builtins.property
  def results_(self) -> _ods_ir.OpResultList:
    _ods_variadic_group_length = len(self.operation.results) - 1 + 1
    return self.operation.results[0:0 + _ods_variadic_group_length]

  @builtins.property
  def alternatives(self) -> _ods_ir.RegionSequence:
    return self.regions[0:]

@_ods_cext.register_op_adaptor(AlternativesOp)
class AlternativesOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "transform.tune.alternatives"

  @builtins.property
  def selected_region_param(self) -> _Optional[_ods_ir.Value]:
    return None if len(self.operands) < 1 else self.operands[0]

  @builtins.property
  def name(self) -> _ods_ir.StringAttr:
    return self.attributes["name"]

  @builtins.property
  def selected_region_attr(self) -> _Optional[_ods_ir.IntegerAttr]:
    if "selected_region_attr" not in self.attributes:
      return None
    return self.attributes["selected_region_attr"]

def tune_alternatives(results_: _Sequence[_ods_ir.Type], name: _Union[_Any, _ods_ir.StringAttr], num_alternatives: int, *, selected_region_attr: _Optional[_Union[_Any, _ods_ir.IntegerAttr]] = None, selected_region_param: _Optional[_ods_ir.Value] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _Union[_ods_ir.OpResult, _ods_ir.OpResultList, AlternativesOp]:
  op = AlternativesOp(results_=results_, name=name, num_alternatives=num_alternatives, selected_region_attr=selected_region_attr, selected_region_param=selected_region_param, loc=loc, ip=ip); results = op.results
  return results if len(results) > 1 else (results[0] if len(results) == 1 else op)

@_ods_cext.register_operation(_Dialect)
class KnobOp(_ods_ir.OpView):
  r"""
  Provides a representation for "tunables" within schedules.
  
  Each op represents a single tunable, which has a `name` and a set
  of valid `options` described by an attribute. Without a specified
  `selected` option, this op represents a non-deterministic choice
  that has yet to be resolved -- as such, the interpreter runtime
  semantics is to raise a failure.
  
  The non-deterministic choice is resolved through providing a
  `selected` attribute. When provided, the interpreter runtime
  semantics are to return the `selected` attribute as a param through
  the op's result.
  
  -----
  
  In case the `options` attribute is an `ArrayAttr`, the verifier
  checks that the provided `selected` attribute occurs in `options`.
  """

  OPERATION_NAME = "transform.tune.knob"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _ods_ir.Type, name: _Union[_Any, _ods_ir.StringAttr], options: _Union[_Any, _ods_ir.Attribute], *, selected: _Optional[_Union[_Any, _ods_ir.Attribute]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["name"] = (name if (
    isinstance(name, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('Builtin_StringAttr')) else
      _ods_ir.AttrBuilder.get('Builtin_StringAttr')(name, context=_ods_context))
    attributes["options"] = (options if (
    isinstance(options, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('AnyAttr')) else
      _ods_ir.AttrBuilder.get('AnyAttr')(options, context=_ods_context))
    if selected is not None: attributes["selected"] = (selected if (
        isinstance(selected, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('AnyAttr')) else
          _ods_ir.AttrBuilder.get('AnyAttr')(selected, context=_ods_context))
    results = []
    results.append(result)
    _ods_successors = None
    super().__init__(self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS, attributes=attributes, results=results, operands=operands, successors=_ods_successors, regions=regions, loc=loc, ip=ip)

  @builtins.property
  def name(self) -> _ods_ir.StringAttr:
    return self.operation.attributes["name"]

  @name.setter
  def name(self, value: _ods_ir.StringAttr):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["name"] = value

  @builtins.property
  def options(self) -> _ods_ir.Attribute:
    return self.operation.attributes["options"]

  @options.setter
  def options(self, value: _ods_ir.Attribute):
    if value is None:
      raise ValueError("'None' not allowed as value for mandatory attributes")
    self.operation.attributes["options"] = value

  @builtins.property
  def selected(self) -> _Optional[_ods_ir.Attribute]:
    if "selected" not in self.operation.attributes:
      return None
    return self.operation.attributes["selected"]

  @selected.setter
  def selected(self, value: _Optional[_ods_ir.Attribute]):
    if value is not None:
      self.operation.attributes["selected"] = value
    elif "selected" in self.operation.attributes:
      del self.operation.attributes["selected"]

  @selected.deleter
  def selected(self):
    del self.operation.attributes["selected"]

  @builtins.property
  def result(self) -> _ods_ir.OpResult:
    return self.operation.results[0]

@_ods_cext.register_op_adaptor(KnobOp)
class KnobOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "transform.tune.knob"

  @builtins.property
  def name(self) -> _ods_ir.StringAttr:
    return self.attributes["name"]

  @builtins.property
  def options(self) -> _ods_ir.Attribute:
    return self.attributes["options"]

  @builtins.property
  def selected(self) -> _Optional[_ods_ir.Attribute]:
    if "selected" not in self.attributes:
      return None
    return self.attributes["selected"]

def tune_knob(result: _ods_ir.Type, name: _Union[_Any, _ods_ir.StringAttr], options: _Union[_Any, _ods_ir.Attribute], *, selected: _Optional[_Union[_Any, _ods_ir.Attribute]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult:
  return KnobOp(result=result, name=name, options=options, selected=selected, loc=loc, ip=ip).result
