[Mosaic] Allow passing ApplyVectorLayoutCtx to tpu.apply_layout_op.

To make it the same with C++ API. While I'm here, fix a bug in test_concatenate. PiperOrigin-RevId: 716016244
jax-ml · Jan 16, 2025 · 4a9cc9f · 4a9cc9f
1 parent d3ba1eb
commit 4a9cc9f
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 12 deletions.
diff --git a/jaxlib/mlir/_mlir_libs/tpu_ext.cc b/jaxlib/mlir/_mlir_libs/tpu_ext.cc
@@ -693,10 +693,9 @@ NB_MODULE(_tpu_ext, m) {
   });
 
   m.def("apply_layout_op",
-        [](int hardware_generation, const MlirOperation c_op) {
+        [](MlirTpuApplyVectorLayoutContext ctx, const MlirOperation c_op) {
           DiagnosticCapture diag_capture(getDefaultContext());
-          MlirLogicalResult res =
-              mlirTpuApplyLayoutOp(hardware_generation, c_op, TARGET_SHAPE);
+          MlirLogicalResult res = mlirTpuApplyLayoutOp(ctx, c_op);
           if (mlirLogicalResultIsFailure(res)) {
             diag_capture.throwIfError();
             throw std::runtime_error("applyLayoutOp failed");

diff --git a/jaxlib/mosaic/dialect/tpu/integrations/c/tpu_dialect.cc b/jaxlib/mosaic/dialect/tpu/integrations/c/tpu_dialect.cc
@@ -383,13 +383,10 @@ MlirTpuValueArray mlirTpuDisassemble(MlirTpuInsertionPoint insertion_point,
   return MlirTpuValueArrayFromXlaArray(std::move(failure_or_vals).value());
 }
 
-MlirLogicalResult mlirTpuApplyLayoutOp(int hardware_generation,
-                                       MlirOperation op,
-                                       MlirTpuI64TargetTuple target_shape) {
-  mlir::tpu::ApplyVectorLayoutContext ctx{
-      .hardware_generation = hardware_generation,
-      .target_shape = unwrap(target_shape)};
-  return wrap(mlir::tpu::applyLayoutOp(ctx, *unwrap(op)));
+MlirLogicalResult mlirTpuApplyLayoutOp(MlirTpuApplyVectorLayoutContext ctx,
+                                       MlirOperation op) {
+  mlir::tpu::ApplyVectorLayoutContext unwrapped_ctx = unwrap(ctx);
+  return wrap(mlir::tpu::applyLayoutOp(unwrapped_ctx, *unwrap(op)));
 }
 
 MlirValue mlirTpuRelayout(MlirTpuInsertionPoint insertion_point, MlirValue val,

diff --git a/jaxlib/mosaic/dialect/tpu/integrations/c/tpu_dialect.h b/jaxlib/mosaic/dialect/tpu/integrations/c/tpu_dialect.h
@@ -227,8 +227,7 @@ MLIR_CAPI_EXPORTED MlirTpuValueArray mlirTpuDisassemble(
     MlirValue val, MlirTpuI64TargetTuple target_shape);
 
 MLIR_CAPI_EXPORTED MlirLogicalResult
-mlirTpuApplyLayoutOp(int hardware_generation, MlirOperation op,
-                     MlirTpuI64TargetTuple target_shape);
+mlirTpuApplyLayoutOp(MlirTpuApplyVectorLayoutContext ctx, MlirOperation op);
 
 // Returns null on failure
 MLIR_CAPI_EXPORTED MlirValue