# Copyright 2024 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.

"""Python bindings for the MLIR Mosaic GPU dialect.

Note: this file *must* be called `mosaic_gpu.py`, in order to match the dialect
name. Otherwise, MLIR is unable to find the module during dialect search.
"""

# ruff: noqa: F401
# ruff: noqa: F403
from jaxlib.mosaic.dialect.gpu._mosaic_gpu_gen_ops import *
from jaxlib.mosaic.dialect.gpu import _mosaic_gpu_gen_ops
from jaxlib.mosaic.dialect.gpu._mosaic_gpu_gen_enums import *
from jaxlib.mlir._mlir_libs._mosaic_gpu_ext import *

try:
  from jaxlib.mlir.dialects._ods_common import _cext
except ImportError:
  from mlir.dialects._ods_common import _cext


# Add the parent module to the search prefix
_cext.globals.append_dialect_search_prefix(__name__[:__name__.rfind(".")])


@_cext.register_operation(_mosaic_gpu_gen_ops._Dialect, replace=True)
class AllocaScopeOp(_mosaic_gpu_gen_ops.AllocaScopeOp):  # noqa: F405
  """An extension to the automatically generated AllocaScopeOp bindings."""

  def __init__(self, result=(), *, loc=None, ip=None):
    super().__init__(result, loc=loc, ip=ip)
    self.regions[0].blocks.append()  # Append the block.

  @property
  def body(self):
    return self.regions[0].blocks[0]


def alloca_scope(result=(), *, loc=None, ip=None) -> AllocaScopeOp:
  return AllocaScopeOp(result, loc=loc, ip=ip)


@_cext.register_operation(_mosaic_gpu_gen_ops._Dialect, replace=True)
class WarpMapOp(_mosaic_gpu_gen_ops.WarpMapOp):  # noqa: F405
  """An extension to the automatically generated WarpMapOp bindings."""

  def __init__(self, operands, *, loc=None, ip=None):
    super().__init__(operands, loc=loc, ip=ip)
    args_ty = [o.type for o in self.operands_]
    self.regions[0].blocks.append(*args_ty)  # Append the block.

  @property
  def body(self):
    return self.regions[0].blocks[0]


def warp_map(operands, *, loc=None, ip=None) -> WarpMapOp:
  return WarpMapOp(operands, loc=loc, ip=ip)


@_cext.register_operation(_mosaic_gpu_gen_ops._Dialect, replace=True)
class RunScopedOp(_mosaic_gpu_gen_ops.RunScopedOp):  # noqa: F405
  """An extension to the automatically generated RunScopedOp bindings."""

  def __init__(
      self,
      buffer_types,
      buffer_attrs,
      *,
      results=(),
      loc=None,
      ip=None,
  ):
    super().__init__(results, buffer_attrs, loc=loc, ip=ip)
    self.regions[0].blocks.append(*buffer_types)  # Append the block.

  @property
  def body(self):
    return self.regions[0].blocks[0]


def run_scoped(
    buffer_types,
    buffer_attrs,
    *,
    results=(),
    loc=None,
    ip=None,
) -> RunScopedOp:
  return RunScopedOp(
      buffer_types,
      buffer_attrs,
      results=results,
      loc=loc,
      ip=ip,
  )