[Mosaic TPU] Allow downgrading the IR during serialization for forward compat

This is to uphold the monthly stability promise made by jax.export.

PiperOrigin-RevId: 704233290

Author: apaszke

jax/_src/tpu_custom_call.py

@@ -62,6 +62,11 @@
     help="Allow hlo dialects in Mosaic",
 )
 
+
+# This tracks the latest Mosaic IR version with a monthly delay.
+FWD_COMPAT_IR_VERSION = 3
+
+
 tpu_custom_call_p = core.Primitive("tpu_custom_call")
 tpu_custom_call_p.def_impl(
     functools.partial(xla.apply_primitive, tpu_custom_call_p))
@@ -407,6 +412,7 @@ def _lower_mosaic_module_to_asm(
     backend: str,
     device_type: str | None,
     kernel_name: str | None,
+    ir_version: int | None = None,
 ) -> tuple[ir.Module, tuple[bool, bool, bool, bool]]:
   has_communication, has_custom_barrier = tpu.private_has_communication(
       module.operation
@@ -438,8 +444,17 @@ def _lower_mosaic_module_to_asm(
       module_op = module.operation.clone()
     prev_allow_unregistered_dialects = ctx.allow_unregistered_dialects
     ctx.allow_unregistered_dialects = True
+    # TODO(apaszke): Remove once the minimum jaxlib version is at least 0.4.37.
+    if jax.version._version_as_tuple(jax.lib.__version__) < (0, 4, 37):
+      target_version = ""
+    else:
+      target_version = (
+          f"target-version={ir_version}" if ir_version is not None else ""
+      )
     try:
-      pipeline = PassManager.parse("builtin.module(mosaic-serde{serialize=true})")
+      pipeline = PassManager.parse(
+          "builtin.module(mosaic-serde{serialize=true " + target_version + "})"
+      )
       pipeline.run(module_op)
     finally:
       ctx.allow_unregistered_dialects = prev_allow_unregistered_dialects
@@ -506,6 +521,7 @@ def _lower_to_custom_call_config(
     serialization_format: int | None,
     output_memory_spaces: tuple[MemorySpace | None, ...] | None = None,
     kernel_name: str | None = None,
+    ir_version: int | None = None,
 ) -> CustomCallBackendConfig:
   lowered_module_asm, (
       has_communication,
@@ -517,6 +533,7 @@ def _lower_to_custom_call_config(
       backend=backend,
       device_type=device_type,
       kernel_name=kernel_name,
+      ir_version=ir_version,
   )
   return _lowered_to_custom_call_config(
       lowered_module_asm,
@@ -617,6 +634,7 @@ def lower_module_to_custom_call(
       serialization_format=serialization_format,
       output_memory_spaces=output_memory_spaces,
       kernel_name=kernel_name,
+      ir_version=FWD_COMPAT_IR_VERSION if ctx.is_forward_compat() else None,
   )
   return _tpu_custom_call_lowering(
       ctx,

jaxlib/mosaic/dialect/tpu/tpu.td

@@ -780,7 +780,10 @@ def DebugAssertInsertionPass : Pass<"debug-assert-insertion", "::mlir::func::Fun
 }
 
 def MosaicSerdePass : Pass<"mosaic-serde", "::mlir::ModuleOp"> {
-  let options = [Option<"serialize", "serialize", "bool", "", "">];
+  let options = [
+    Option<"serialize", "serialize", "bool", "", "">,
+    Option<"target_version", "target-version", "int", "", "">  // Only used when serialize=true.
+  ];
 }
 
 def LogicalToPhysicalDeviceIdPass : Pass<"logical-to-physical-device-id", "::mlir::func::FuncOp"> {

jaxlib/mosaic/dialect/tpu/transforms/serde.cc

@@ -43,6 +43,8 @@ namespace {
 
 constexpr std::string_view kMangledDialect = "stable_mosaic.";
 constexpr StringRef kVersionAttrName = "stable_mosaic.version";
+// When this is bumped, we should file a TODO to update the forward-compatible
+// version in tpu_custom_call.py in a month!
 constexpr int kVersion = 3;
 
 StringRef mangle(StringRef name, std::string* storage) {
@@ -63,7 +65,7 @@ std::optional<StringRef> demangle(StringRef name) {
 
 using rule_type = std::function<LogicalResult(Operation*, int)>;
 
-LogicalResult enqueue_dma_rule(Operation* op, int version) {
+LogicalResult enqueue_dma_upgrade(Operation* op, int version) {
   // Added AttrSizedOperandSegments and core_id in version 2.
   if (version < 2) {
     if (op->getNumOperands() == 3) {  // Local DMA.
@@ -84,17 +86,21 @@ LogicalResult enqueue_dma_rule(Operation* op, int version) {
   return success();
 }
 
-LogicalResult semaphore_signal_rule(Operation* op, int version) {
+LogicalResult enqueue_dma_downgrade(Operation* op, int version) {
+  if (version < 2) {
+    return op->emitError("Downgrade to version ") << version << " unsupported";
+  }
+  return success();
+}
+
+LogicalResult semaphore_signal_upgrade(Operation* op, int version) {
   // Added AttrSizedOperandSegments and core_id in version 2.
   if (version < 2) {
     if (op->getNumOperands() == 2) {  // Local signal.
       op->setAttr(OpTrait::AttrSizedOperandSegments<
                       EnqueueDMAOp>::getOperandSegmentSizeAttr(),
                   mlir::DenseI32ArrayAttr::get(op->getContext(), {1, 1, 0, 0}));
     } else if (op->getNumOperands() == 3) {  // Remote signal.
-      // Hardcoding that one optional value is device_id, not core_id. This
-      // could misinterpret sem_signals where core_id is specified, but
-      // device_id isn't.
       op->setAttr(OpTrait::AttrSizedOperandSegments<
                       EnqueueDMAOp>::getOperandSegmentSizeAttr(),
                   mlir::DenseI32ArrayAttr::get(op->getContext(), {1, 1, 1, 0}));
@@ -105,7 +111,25 @@ LogicalResult semaphore_signal_rule(Operation* op, int version) {
   return success();
 }
 
-LogicalResult vector_multi_dim_reduce_rule(Operation* op, int version) {
+LogicalResult semaphore_signal_downgrade(Operation* op, int version) {
+  if (version < 2) {
+    auto operands = op->getAttrOfType<mlir::DenseI32ArrayAttr>(
+        OpTrait::AttrSizedOperandSegments<
+            EnqueueDMAOp>::getOperandSegmentSizeAttr());
+    if (!operands || operands.size() != 4) {
+      return op->emitError("Missing or invalid AttrSizedOperandSegments");
+    }
+    if (operands[3]) {
+      return op->emitError("Downgrade to version ")
+             << version << " impossible: core_id is set";
+    }
+    op->removeAttr(OpTrait::AttrSizedOperandSegments<
+                   EnqueueDMAOp>::getOperandSegmentSizeAttr());
+  }
+  return success();
+}
+
+LogicalResult vector_multi_dim_reduce_upgrade(Operation* op, int version) {
   // Changed reductions_dims from ArrayAttr of IntegerAttrs to DenseI64ArrayAttr
   // in version 3.
   if (version < 3) {
@@ -133,21 +157,49 @@ LogicalResult vector_multi_dim_reduce_rule(Operation* op, int version) {
   return success();
 }
 
+LogicalResult vector_multi_dim_reduce_downgrade(Operation* op, int version) {
+  if (version < 3) {
+    return op->emitError("Downgrade to version ") << version << " unsupported";
+  }
+  return success();
+}
+
 const llvm::StringMap<rule_type>& upgrade_rules() {
   static auto rules = new llvm::StringMap<rule_type>{
-      {EnqueueDMAOp::getOperationName(), enqueue_dma_rule},
-      {SemaphoreSignalOp::getOperationName(), semaphore_signal_rule},
+      {EnqueueDMAOp::getOperationName(), enqueue_dma_upgrade},
+      {SemaphoreSignalOp::getOperationName(), semaphore_signal_upgrade},
       {vector::MultiDimReductionOp::getOperationName(),
-       vector_multi_dim_reduce_rule}
+       vector_multi_dim_reduce_upgrade}
   };
   return *rules;
 }
 
+const llvm::StringMap<rule_type>& downgrade_rules() {
+  static auto rules = new llvm::StringMap<rule_type>{
+      {EnqueueDMAOp::getOperationName(), enqueue_dma_downgrade},
+      {SemaphoreSignalOp::getOperationName(), semaphore_signal_downgrade},
+      {vector::MultiDimReductionOp::getOperationName(),
+       vector_multi_dim_reduce_downgrade}};
+  return *rules;
+}
+
 struct MosaicSerdePass : public impl::MosaicSerdePassBase<MosaicSerdePass> {
   using Base::Base;
 
   void runOnOperation() override {
     ModuleOp module = getOperation();
+    if (!serialize.hasValue()) {
+      module.emitError("serialize option must be specified");
+      return signalPassFailure();
+    }
+    int serialize_version =
+        target_version.hasValue() ? target_version : kVersion;
+    if (serialize && serialize_version > kVersion) {
+      module.emitError("The highest supported version is ")
+          << kVersion << " but requested serialization at version "
+          << serialize_version;
+      return signalPassFailure();
+    }
     if (serialize && !module->getContext()->allowsUnregisteredDialects()) {
       module.emitError() << "Cannot serialize within a context that does not "
                             "allow unregistered dialects.";
@@ -159,7 +211,7 @@ struct MosaicSerdePass : public impl::MosaicSerdePassBase<MosaicSerdePass> {
       module->setAttr(
           kVersionAttrName,
           IntegerAttr::get(IntegerType::get(module->getContext(), 64),
-                           kVersion));
+                           serialize_version));
     } else {
       IntegerAttr version_attr =
           module->getAttrOfType<IntegerAttr>(kVersionAttrName);
@@ -178,7 +230,7 @@ struct MosaicSerdePass : public impl::MosaicSerdePassBase<MosaicSerdePass> {
       module->removeAttr(kVersionAttrName);
     }
     std::string name_storage;
-    auto result = module.walk([this, &name_storage, version](Operation* op) {
+    auto result = module.walk([&](Operation* op) {
       if (isa<ModuleOp>(op)) {  // Don't mangle the ModuleOp itself.
         return WalkResult::advance();
       }
@@ -210,6 +262,16 @@ struct MosaicSerdePass : public impl::MosaicSerdePassBase<MosaicSerdePass> {
       auto new_op = Operation::create(
           op->getLoc(), *new_name, op->getResultTypes(), op->getOperands(),
           op->getAttrs(), nullptr, op->getSuccessors(), op->getRegions());
+      // Downgrade the op to the target version, if needed.
+      if (serialize && kVersion != serialize_version) {
+        if (const auto rule =
+                downgrade_rules().find(op->getName().getStringRef());
+            rule != downgrade_rules().end()) {
+          if (rule->second(new_op, serialize_version).failed()) {
+            return WalkResult::interrupt();
+          }
+        }
+      }
       op->getBlock()->getOperations().insertAfter(Block::iterator(op), new_op);
       op->replaceAllUsesWith(new_op->getResults());
       op->erase();