mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Mosaic TPU][NFC] Create local namespace to prevent function name duplication error under global namespace mlir::tpu

Browse Source 
PiperOrigin-RevId: 705538965
This commit is contained in:
Jevin Jiang 2024-12-12 09:52:58 -08:00 committed by jax authors
parent 3c649b134a
commit 3ff5706051
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
jaxlib/mosaic/dialect/tpu/transforms/canonicalize_mosaic.cc
View File

@ -44,6 +44,8 @@ namespace mlir::tpu {
#define GEN_PASS_DEF_CANONICALIZEMOSAICPASS
#include "jaxlib/mosaic/dialect/tpu/tpu_passes.h.inc"
namespace {
LogicalResult tpu_matmul_rule(tpu::MatmulOp op) {
ImplicitLocOpBuilder builder(op.getLoc(), op.getOperation());
@ -611,6 +613,8 @@ struct CanonicalizeMosaicPass
};
};
} // namespace
std::unique_ptr<OperationPass<func::FuncOp>> createCanonicalizeMosaicPass(
int hardware_generation) {
return std::make_unique<CanonicalizeMosaicPass>(hardware_generation);