
# 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 ApplyNVGPUToNVVMConversionPatternsOp(_ods_ir.OpView):
  r"""
  Collects patterns that convert NVGPU dialect ops to NVVM dialect ops. These
  patterns require an "LLVMTypeConverter".
  """

  OPERATION_NAME = "transform.apply_conversion_patterns.nvgpu.nvgpu_to_nvvm"

  _ODS_REGIONS = (0, True)

  def __init__(self, *, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    _ods_context = _ods_get_default_loc_context(loc)
    results = []
    _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)

@_ods_cext.register_op_adaptor(ApplyNVGPUToNVVMConversionPatternsOp)
class ApplyNVGPUToNVVMConversionPatternsOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "transform.apply_conversion_patterns.nvgpu.nvgpu_to_nvvm"

def apply_conversion_patterns_nvgpu_nvgpu_to_nvvm(*, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> ApplyNVGPUToNVVMConversionPatternsOp:
  return ApplyNVGPUToNVVMConversionPatternsOp(loc=loc, ip=ip)

@_ods_cext.register_operation(_Dialect)
class CreateAsyncGroupsOp(_ods_ir.OpView):
  r"""
  Look for global to shared memory copies within the targeted op in the form
  of vector transfer ops and convert them to async copies when possible.
  Consecutive copies are put into the same group. A "wait" operation is
  inserted right at the of end the group.
  
  `bypass_l1` specifies whether `bypassL1` attributes should be added to
  the async copies. `bypass_l1` is a compiler hint: only 16 byte transfers
  can bypass the L1 cache, so this attribute is not set for any other transfer
  sizes.
  
  #### Return modes
  
  This op consumes the `target` handle and produces the `result` handle, which
  is mapped to the same payload operations as the `target` handle. The op
  modifies the payload.
  """

  OPERATION_NAME = "transform.nvgpu.create_async_groups"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _ods_ir.Type, target: _ods_ir.Value, *, bypass_l1: _Optional[bool] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(target)
    _ods_context = _ods_get_default_loc_context(loc)
    if bool(bypass_l1): attributes["bypass_l1"] = _ods_ir.UnitAttr.get(
      _ods_get_default_loc_context(loc))
    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 target(self) -> _ods_ir.Value:
    return self.operation.operands[0]

  @builtins.property
  def bypass_l1(self) -> bool:
    return "bypass_l1" in self.operation.attributes

  @bypass_l1.setter
  def bypass_l1(self, value):
    if bool(value):
      self.operation.attributes["bypass_l1"] = _ods_ir.UnitAttr.get()
    elif "bypass_l1" in self.operation.attributes:
      del self.operation.attributes["bypass_l1"]

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

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

@_ods_cext.register_op_adaptor(CreateAsyncGroupsOp)
class CreateAsyncGroupsOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "transform.nvgpu.create_async_groups"

  @builtins.property
  def target(self) -> _ods_ir.Value:
    return self.operands[0]

  @builtins.property
  def bypass_l1(self) -> bool:
    return "bypass_l1" in self.attributes

def nvgpu_create_async_groups(result: _ods_ir.Type, target: _ods_ir.Value, *, bypass_l1: _Optional[bool] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult:
  return CreateAsyncGroupsOp(result=result, target=target, bypass_l1=bypass_l1, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class PipelineSharedMemoryCopiesOp(_ods_ir.OpView):
  r"""
  Applies software pipelining to a given scf.for loop. The pipelining
  strategy will look for a load into shared memory and pipeline it to overlap
  it with the rest of the loop.
  
  NOTE: It is user responsibility to ensure that there are no dependency
  between `depth` iterations of the loop by using multi-buffering. It is
  also user responsibility to ensure a sufficient amount of shared memory
  is allocated to cover eventual writes by `depth-1` speculative
  iterations.
  
  `depth` will indicate how many stages the software pipeline should have.
  `peel_epilogue` allows to force the epilogue to be peeled out instead of
  potentially using predicated operations for the epilogue phase.
  
  #### Return modes
  
  Consumes the operand handle and produces a result handle pointing to the
  loop, which may or may not have been pipelined. Produces a definite failure
  if the loop pipeliner mutated the IR before failing to pipeline, in
  particular if `peel_epilogue` is not set and the loop body doesn't support
  predication. If failure propagation mode is set to "propagate", produces a
  silenceable failure when pipelining preconditions, e.g., loop bound being
  static, are not met or when the loop wasn't pipelined because due to the
  lack of loads into shared memory. If the failure propagation mode is set
  to "suppress" (default), succeeds in these case and associates the result
  handle with the original loop.
  
  TODO: the shared memory part and behavior specific to NVGPU should be
  made orthogonal to pipelining so that `transform.loop.pipeline` becomes
  usable here.
  """

  OPERATION_NAME = "transform.nvgpu.pipeline_shared_memory_copies"

  _ODS_REGIONS = (0, True)

  def __init__(self, result: _ods_ir.Type, for_op: _ods_ir.Value, depth: _Union[int, _ods_ir.IntegerAttr], *, peel_epilogue: _Optional[bool] = None, failure_propagation_mode: _Optional[_Union[_Any, _ods_ir.Attribute]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(for_op)
    _ods_context = _ods_get_default_loc_context(loc)
    attributes["depth"] = (depth if (
    isinstance(depth, _ods_ir.Attribute) or
    not _ods_ir.AttrBuilder.contains('I64Attr')) else
      _ods_ir.AttrBuilder.get('I64Attr')(depth, context=_ods_context))
    if bool(peel_epilogue): attributes["peel_epilogue"] = _ods_ir.UnitAttr.get(
      _ods_get_default_loc_context(loc))
    if failure_propagation_mode is not None: attributes["failure_propagation_mode"] = (failure_propagation_mode if (
        isinstance(failure_propagation_mode, _ods_ir.Attribute) or
        not _ods_ir.AttrBuilder.contains('FailurePropagationMode')) else
          _ods_ir.AttrBuilder.get('FailurePropagationMode')(failure_propagation_mode, 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 for_op(self) -> _ods_ir.Value:
    return self.operation.operands[0]

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

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

  @builtins.property
  def peel_epilogue(self) -> bool:
    return "peel_epilogue" in self.operation.attributes

  @peel_epilogue.setter
  def peel_epilogue(self, value):
    if bool(value):
      self.operation.attributes["peel_epilogue"] = _ods_ir.UnitAttr.get()
    elif "peel_epilogue" in self.operation.attributes:
      del self.operation.attributes["peel_epilogue"]

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

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

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

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

@_ods_cext.register_op_adaptor(PipelineSharedMemoryCopiesOp)
class PipelineSharedMemoryCopiesOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "transform.nvgpu.pipeline_shared_memory_copies"

  @builtins.property
  def for_op(self) -> _ods_ir.Value:
    return self.operands[0]

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

  @builtins.property
  def peel_epilogue(self) -> bool:
    return "peel_epilogue" in self.attributes

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

def nvgpu_pipeline_shared_memory_copies(result: _ods_ir.Type, for_op: _ods_ir.Value, depth: _Union[int, _ods_ir.IntegerAttr], *, peel_epilogue: _Optional[bool] = None, failure_propagation_mode: _Optional[_Union[_Any, _ods_ir.Attribute]] = None, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> _ods_ir.OpResult:
  return PipelineSharedMemoryCopiesOp(result=result, for_op=for_op, depth=depth, peel_epilogue=peel_epilogue, failure_propagation_mode=failure_propagation_mode, loc=loc, ip=ip).result

@_ods_cext.register_operation(_Dialect)
class RewriteCopyAsTmaOp(_ods_ir.OpView):
  r"""
  Rewrite a copy operation on memref to tma operations that transit through
  shared memory.
  """

  OPERATION_NAME = "transform.nvgpu.rewrite_copy_as_tma"

  _ODS_REGIONS = (0, True)

  def __init__(self, target: _ods_ir.Value, *, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(target)
    _ods_context = _ods_get_default_loc_context(loc)
    results = []
    _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 target(self) -> _ods_ir.Value:
    return self.operation.operands[0]

@_ods_cext.register_op_adaptor(RewriteCopyAsTmaOp)
class RewriteCopyAsTmaOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "transform.nvgpu.rewrite_copy_as_tma"

  @builtins.property
  def target(self) -> _ods_ir.Value:
    return self.operands[0]

def nvgpu_rewrite_copy_as_tma(target: _ods_ir.Value, *, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> RewriteCopyAsTmaOp:
  return RewriteCopyAsTmaOp(target=target, loc=loc, ip=ip)

@_ods_cext.register_operation(_Dialect)
class RewriteMatmulAsMmaSyncOp(_ods_ir.OpView):
  r"""
  Rewrite a matmul operation on memref to an mma.sync operation on vectors.
  
  Memory copies with the required access patterns are automatically inserted.
  Operations that do not have a 1-1 mapping to mma.sync operations are left
  unchanged.
  """

  OPERATION_NAME = "transform.nvgpu.rewrite_matmul_as_mma_sync"

  _ODS_REGIONS = (0, True)

  def __init__(self, target: _ods_ir.Value, *, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None):
    operands = []
    attributes = {}
    regions = None
    operands.append(target)
    _ods_context = _ods_get_default_loc_context(loc)
    results = []
    _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 target(self) -> _ods_ir.Value:
    return self.operation.operands[0]

@_ods_cext.register_op_adaptor(RewriteMatmulAsMmaSyncOp)
class RewriteMatmulAsMmaSyncOpAdaptor(_ods_ir.OpAdaptor):
  OPERATION_NAME = "transform.nvgpu.rewrite_matmul_as_mma_sync"

  @builtins.property
  def target(self) -> _ods_ir.Value:
    return self.operands[0]

def nvgpu_rewrite_matmul_as_mma_sync(target: _ods_ir.Value, *, loc: _Optional[_ods_ir.Location] = None, ip: _Optional[_ods_ir.InsertionPoint] = None) -> RewriteMatmulAsMmaSyncOp:
  return RewriteMatmulAsMmaSyncOp(target=target, loc=loc, ip=ip)
