[Mosaic GPU] Clean up the module structure

jax/_src/pallas/mosaic_gpu/core.py

@@ -23,7 +23,7 @@
 from jax._src import dtypes
 from jax._src import tree_util
 from jax._src.pallas import core as pallas_core
-from jax.experimental.mosaic import gpu as mosaic_gpu
+import jax.experimental.mosaic.gpu as mgpu
 import jax.numpy as jnp
 
 
@@ -64,7 +64,7 @@ def __call__(self, shape: tuple[int, ...], dtype: jnp.dtype):
 
 
 class MemoryRefTransform(pallas_core.MemoryRefTransform, Protocol):
-  def to_gpu_transform(self) -> mosaic_gpu.MemRefTransform:
+  def to_gpu_transform(self) -> mgpu.MemRefTransform:
     ...
 
 
@@ -101,8 +101,8 @@ def __call__(
         inner_aval=block_aval.inner_aval.update(shape=new_block_shape)
     )
 
-  def to_gpu_transform(self) -> mosaic_gpu.MemRefTransform:
-    return mosaic_gpu.TileTransform(self.tiling)
+  def to_gpu_transform(self) -> mgpu.MemRefTransform:
+    return mgpu.TileTransform(self.tiling)
 
 
 @dataclasses.dataclass(frozen=True)

jax/_src/pallas/mosaic_gpu/lowering.py

@@ -39,8 +39,7 @@
 from jax._src.pallas import utils as pallas_utils
 from jax._src.pallas.mosaic_gpu import core as gpu_core
 from jax._src.state import primitives as sp
-from jax.experimental.mosaic import gpu as mosaic_gpu
-from jax.experimental.mosaic.gpu import dsl as mgpu
+import jax.experimental.mosaic.gpu as mgpu
 import jax.numpy as jnp
 import numpy as np
 
@@ -160,7 +159,7 @@ def stack_free_smem(self, bytes: int):
 @dataclasses.dataclass(frozen=True)
 class LoweringRuleContext:
   module_ctx: ModuleContext
-  launch_ctx: mosaic_gpu.LaunchContext
+  launch_ctx: mgpu.LaunchContext
   avals_in: Sequence[jax_core.ShapedArray]
   avals_out: Sequence[jax_core.ShapedArray]
 
@@ -180,7 +179,7 @@ class LoweringError(Exception):  # pylint: disable=g-bad-exception-name
 
 def _eval_index_map(
     module_ctx: ModuleContext,
-    launch_ctx: mosaic_gpu.LaunchContext,
+    launch_ctx: mgpu.LaunchContext,
     idx: ir.Value,
     block_mapping: pallas_core.BlockMapping,
 ) -> Sequence[ir.Value]:
@@ -300,7 +299,7 @@ def lower_jaxpr_to_module(
       )
   ]
 
-  def body(launch_ctx: mosaic_gpu.LaunchContext, *buffers: ir.Value):
+  def body(launch_ctx: mgpu.LaunchContext, *buffers: ir.Value):
     *buffers_gmem, (
         buffers_smem,
         *scratch_buffers_smem,
@@ -494,7 +493,7 @@ def _(step, _):
       jax.ShapeDtypeStruct(shape=[smem_scratch_bytes], dtype=np.int8)
   )
 
-  module, out_structs_smem, _ = mosaic_gpu._lower_as_gpu_kernel(
+  module, out_structs_smem, _ = mgpu._lower_as_gpu_kernel(
       body,
       grid=grid,
       cluster=(),
@@ -528,7 +527,7 @@ def deco(fn):
 
 def lower_jaxpr_to_mosaic_gpu(
     module_ctx: ModuleContext,
-    launch_ctx: mosaic_gpu.LaunchContext,
+    launch_ctx: mgpu.LaunchContext,
     jaxpr: jax_core.Jaxpr,
     args: Sequence[ir.Value],
     consts=(),

jax/_src/pallas/mosaic_gpu/pallas_call_registration.py

@@ -23,7 +23,7 @@
 from jax._src.interpreters import mlir
 from jax._src.pallas import core as pallas_core
 from jax._src.pallas.mosaic_gpu import lowering
-from jax.experimental.mosaic import gpu as mosaic_gpu
+import jax.experimental.mosaic.gpu.core as mosaic_core
 
 
 def pallas_call_lowering(
@@ -67,7 +67,7 @@ def pallas_call_lowering(
     print(lowering_result.module.operation)
 
   module = lowering_result.module
-  return mosaic_gpu._mosaic_gpu_lowering_rule(
+  return mosaic_core._mosaic_gpu_lowering_rule(
       ctx,
       *args,
       module=module.operation.get_asm(binary=True, enable_debug_info=True),