hand

venv/lib/python3.12/site-packages/jaxlib/mosaic/python/mosaic_gpu.py

@@ -0,0 +1,53 @@
+# 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 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)