# Copyright 2026 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# THIS FILE IS AUTOMATICALLY GENERATED. DO NOT EDIT.

from collections.abc import Sequence
from typing import Final

from mlir import ir

def register_dialect(context: ir.Context, load: bool = ...) -> None: ...
def register_inliner_extensions(arg: ir.Context, /) -> None: ...

class DumpOptions:
  def __init__(self) -> None: ...
  @property
  def constraint_system(self) -> bool: ...
  @property
  def mlir_passes(self) -> bool: ...
  @property
  def ptx(self) -> bool: ...
  @property
  def ptxas(self) -> bool: ...
  @property
  def resources(self) -> bool: ...
  @property
  def sass(self) -> bool: ...
  @property
  def sass_ctrl(self) -> bool: ...
  @property
  def dump_path(self) -> str: ...
  @property
  def module_basename(self) -> str: ...

def get_or_set_dump_options(module: ir.Module) -> DumpOptions: ...

class BarrierType(ir.Type):
  def __init__(self, cast_from_type: ir.Type) -> None: ...

  static_typeid: Final[ir.TypeID] = ...
  """(arg: object, /) -> ir.TypeID"""

  @property
  def typeid(self) -> ir.TypeID: ...
  def __repr__(self) -> str: ...
  @staticmethod
  def get(
      orders_tensor_core: bool = ..., ctx: ir.Context | None = ...
  ) -> BarrierType: ...
  @property
  def orders_tensor_core(self) -> bool: ...

class B6x16P32Type(ir.Type):
  def __init__(self, cast_from_type: ir.Type) -> None: ...

  static_typeid: Final[ir.TypeID] = ...
  """(arg: object, /) -> ir.TypeID"""

  @property
  def typeid(self) -> ir.TypeID: ...
  def __repr__(self) -> str: ...
  @staticmethod
  def get(
      element_type: ir.Type, ctx: ir.Context | None = ...
  ) -> B6x16P32Type: ...
  @property
  def element_type(self) -> ir.Type: ...

class P2B6Type(ir.Type):
  def __init__(self, cast_from_type: ir.Type) -> None: ...

  static_typeid: Final[ir.TypeID] = ...
  """(arg: object, /) -> ir.TypeID"""

  @property
  def typeid(self) -> ir.TypeID: ...
  def __repr__(self) -> str: ...
  @staticmethod
  def get(element_type: ir.Type, ctx: ir.Context | None = ...) -> P2B6Type: ...
  @property
  def element_type(self) -> ir.Type: ...

class TileTransformAttr(ir.Attribute):
  def __init__(self, cast_from_attr: ir.Attribute) -> None: ...
  @property
  def type(self) -> ir.Type: ...

  static_typeid: Final[ir.TypeID] = ...
  """(arg: object, /) -> ir.TypeID"""

  @property
  def typeid(self) -> ir.TypeID: ...
  def __repr__(self) -> str: ...
  @staticmethod
  def get(
      tiling: Sequence[int], ctx: ir.Context | None = ...
  ) -> TileTransformAttr: ...
  @property
  def tiling(self) -> ir.DenseI32ArrayAttr: ...

class SwizzleTransformAttr(ir.Attribute):
  def __init__(self, cast_from_attr: ir.Attribute) -> None: ...
  @property
  def type(self) -> ir.Type: ...

  static_typeid: Final[ir.TypeID] = ...
  """(arg: object, /) -> ir.TypeID"""

  @property
  def typeid(self) -> ir.TypeID: ...
  def __repr__(self) -> str: ...
  @staticmethod
  def get(
      swizzle: int, ctx: ir.Context | None = ...
  ) -> SwizzleTransformAttr: ...
  @property
  def swizzle(self) -> int: ...

class WGSplatFragLayoutAttr(ir.Attribute):
  def __init__(self, cast_from_attr: ir.Attribute) -> None: ...
  @property
  def type(self) -> ir.Type: ...

  static_typeid: Final[ir.TypeID] = ...
  """(arg: object, /) -> ir.TypeID"""

  @property
  def typeid(self) -> ir.TypeID: ...
  def __repr__(self) -> str: ...
  @staticmethod
  def get(
      shape: ir.DenseI64ArrayAttr, ctx: ir.Context | None = None
  ) -> WGSplatFragLayoutAttr: ...
  @property
  def shape(self) -> ir.DenseI64ArrayAttr: ...

class WGStridedFragLayoutAttr(ir.Attribute):
  def __init__(self, cast_from_attr: ir.Attribute) -> None: ...
  @property
  def type(self) -> ir.Type: ...

  static_typeid: Final[ir.TypeID] = ...
  """(arg: object, /) -> ir.TypeID"""

  @property
  def typeid(self) -> ir.TypeID: ...
  def __repr__(self) -> str: ...
  @staticmethod
  def get(
      shape: ir.DenseI64ArrayAttr,
      vector_size: int,
      ctx: ir.Context | None = None,
  ) -> WGStridedFragLayoutAttr: ...
  @property
  def shape(self) -> ir.DenseI64ArrayAttr: ...
  @property
  def vector_size(self) -> int: ...

class ReplicatedAttr(ir.Attribute):
  def __init__(self, cast_from_attr: ir.Attribute) -> None: ...
  @property
  def type(self) -> ir.Type: ...

  static_typeid: Final[ir.TypeID] = ...
  """(arg: object, /) -> ir.TypeID"""

  @property
  def typeid(self) -> ir.TypeID: ...
  def __repr__(self) -> str: ...
  @staticmethod
  def get(times: int, ctx: ir.Context | None = ...) -> ReplicatedAttr: ...
  @property
  def times(self) -> int: ...

class TiledLayoutAttr(ir.Attribute):
  def __init__(self, cast_from_attr: ir.Attribute) -> None: ...
  @property
  def type(self) -> ir.Type: ...

  static_typeid: Final[ir.TypeID] = ...
  """(arg: object, /) -> ir.TypeID"""

  @property
  def typeid(self) -> ir.TypeID: ...
  def __repr__(self) -> str: ...
  @staticmethod
  def get(
      tiling: ir.ArrayAttr,
      warp_dims: ir.ArrayAttr,
      lane_dims: ir.ArrayAttr,
      vector_dim: int,
      ctx: ir.Context | None = None,
  ) -> TiledLayoutAttr: ...
  @property
  def tiling(self) -> ir.ArrayAttr: ...
  @property
  def warp_dims(self) -> ir.ArrayAttr: ...
  @property
  def lane_dims(self) -> ir.ArrayAttr: ...
  @property
  def vector_dim(self) -> int: ...

class CopyPartitionAttrInterface(ir.Attribute):
  def __init__(self, cast_from_attr: ir.Attribute) -> None: ...
  @property
  def type(self) -> ir.Type: ...

  static_typeid: Final[ir.TypeID] = ...
  """(arg: object, /) -> ir.TypeID"""

  @property
  def typeid(self) -> ir.TypeID: ...
  def __repr__(self) -> str: ...

class CopyReplicatedAttr(CopyPartitionAttrInterface):
  def __init__(self, cast_from_attr: ir.Attribute) -> None: ...
  @property
  def type(self) -> ir.Type: ...
  @property
  def typeid(self) -> ir.TypeID: ...
  def __repr__(self) -> str: ...
  @staticmethod
  def get(ctx: ir.Context | None = ...) -> CopyReplicatedAttr: ...

class CopyPartitionedAttr(CopyPartitionAttrInterface):
  def __init__(self, cast_from_attr: ir.Attribute) -> None: ...
  @property
  def type(self) -> ir.Type: ...
  @property
  def typeid(self) -> ir.TypeID: ...
  def __repr__(self) -> str: ...
  @staticmethod
  def get(axis: int, ctx: ir.Context | None = ...) -> CopyPartitionedAttr: ...
  @property
  def axis(self) -> int: ...

class SmemAllocAttr(ir.Attribute):
  def __init__(self, cast_from_attr: ir.Attribute) -> None: ...
  @property
  def type(self) -> ir.Type: ...

  static_typeid: Final[ir.TypeID] = ...
  """(arg: object, /) -> ir.TypeID"""

  @property
  def typeid(self) -> ir.TypeID: ...
  def __repr__(self) -> str: ...
  @staticmethod
  def get(
      alignment: int | None = ..., ctx: ir.Context | None = ...
  ) -> SmemAllocAttr: ...
  @property
  def alignment(self) -> int | None: ...