[mosaic] infer-memref-layout is now defined via pass_boilerplate.h

There is little benefit in having pass definitions in the .td file.

PiperOrigin-RevId: 834373569

Author: superbobry

jaxlib/mosaic/dialect/tpu/tpu.td

@@ -1489,22 +1489,6 @@ def LogicalToPhysicalDeviceIdPass : Pass<"logical-to-physical-device-id", "::mli
   let options = [Option<"total_devices", "total-devices", "int", "", "">];
 }
 
-def InferMemRefLayoutPass : Pass<"tpu-infer-memref-layout", "::mlir::func::FuncOp"> {
-  let dependentDialects = [
-    "::mlir::func::FuncDialect",
-    "::mlir::memref::MemRefDialect",
-  ];
-  let constructor = "::mlir::tpu::createInferMemRefLayoutPass()";
-  let options = [
-    // If hardware_generation is not set, the default value of -1 will crash on
-    // runOnOperation.
-    Option<"hardware_generation", "hardware-generation", "int", /*default=*/"-1", "">,
-    Option<"lane_count", "lane-count", "int", /*default=*/"128", "">,
-    Option<"sublane_count", "sublane-count", "int", /*default=*/"8", "">,
-    Option<"tpu_tiling_flags", "tpu-tiling-flags", "::mlir::tpu::TpuTilingFlags", /*default=*/"::mlir::tpu::TpuTilingFlags{}", "">,
-  ];
-}
-
 def CanonicalizeMosaicPass : Pass<"tpu-canonicalize-mosaic", "::mlir::func::FuncOp"> {
   let dependentDialects = [
     "::mlir::arith::ArithDialect",

jaxlib/mosaic/dialect/tpu/tpu_dialect.h

@@ -70,9 +70,8 @@ struct ApplyVectorLayoutContext {
 std::pair<bool, bool> mightCommunicateBetweenChips(Operation *op);
 
 std::unique_ptr<OperationPass<func::FuncOp>> createInferMemRefLayoutPass(
-    int hardware_generation = -1,
-    std::array<int64_t, 2> target_shape = {8, 128},
-    const TpuTilingFlags &tpu_tiling_flags = {});
+    int hardware_generation, std::array<int64_t, 2> target_shape,
+    const TpuTilingFlags& tpu_tiling_flags);
 
 std::unique_ptr<OperationPass<func::FuncOp>> createCanonicalizeMosaicPass(
     int hardware_generation = -1, bool compatibility_mode = true,