From 2948853eea07dba3df54eb2533495c3f40a03a3d Mon Sep 17 00:00:00 2001 From: jax authors Date: Thu, 2 Oct 2025 22:39:36 -0700 Subject: [PATCH] Support For Roofline Analysis Of Pallas Kernels PiperOrigin-RevId: 814542103 --- jaxlib/mosaic/BUILD | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/jaxlib/mosaic/BUILD b/jaxlib/mosaic/BUILD index 9f61520510cc..ec0256c51a1b 100644 --- a/jaxlib/mosaic/BUILD +++ b/jaxlib/mosaic/BUILD @@ -60,7 +60,10 @@ cc_library( ] + glob([ "dialect/tpu/transforms/*.h", ]), - # compatible with libtpu + compatible_with = [ + "//buildenv/target:libtpu", + "//buildenv/target:non_prod", + ], deps = [ ":pass_boilerplate", ":serde", @@ -104,7 +107,10 @@ cc_library( gentbl_cc_library( name = "tpu_inc_gen", - # compatible with libtpu + compatible_with = [ + "//buildenv/target:libtpu", + "//buildenv/target:non_prod", + ], tbl_outs = { "dialect/tpu/tpu_ops.h.inc": ["-gen-op-decls"], "dialect/tpu/tpu_ops.cc.inc": ["-gen-op-defs"], @@ -139,7 +145,10 @@ td_library( srcs = [ "dialect/tpu/tpu.td", ], - # compatible with libtpu + compatible_with = [ + "//buildenv/target:libtpu", + "//buildenv/target:non_prod", + ], deps = [ "@llvm-project//mlir:BuiltinDialectTdFiles", "@llvm-project//mlir:ControlFlowInterfacesTdFiles", @@ -263,7 +272,10 @@ filegroup( cc_library( name = "pass_boilerplate", hdrs = ["pass_boilerplate.h"], - # compatible with libtpu + compatible_with = [ + "//buildenv/target:libtpu", + "//buildenv/target:non_prod", + ], deps = [ "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass", @@ -275,7 +287,10 @@ cc_library( name = "serde", srcs = ["serde.cc"], hdrs = ["serde.h"], - # compatible with libtpu + compatible_with = [ + "//buildenv/target:libtpu", + "//buildenv/target:non_prod", + ], deps = [ "@llvm-project//llvm:Support", "@llvm-project//mlir:DataLayoutInterfaces",