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

Integrate LLVM at llvm/llvm-project@912b154f3a

Browse Source 
Updates LLVM usage to match
[912b154f3a3f](https://github.com/llvm/llvm-project/commit/912b154f3a3f)

PiperOrigin-RevId: 727895384
This commit is contained in:
jax authors 2025-02-17 10:07:50 -08:00
parent bffeb5d60e
commit e78a469b42
2 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
jaxlib/mosaic/gpu/BUILD
View File

@ -168,6 +168,7 @@ cc_library(
"@llvm-project//mlir:Transforms",
"@llvm-project//mlir:UBToLLVM",
"@llvm-project//mlir:VectorDialect",
"@llvm-project//mlir:VectorToLLVM",
"@xla//xla/service:custom_call_status",
"@xla//xla/service:custom_call_target_registry",
],

2
jaxlib/mosaic/gpu/custom_call.cc
View File

@ -53,6 +53,7 @@ limitations under the License.
#include "mlir/include/mlir/Conversion/NVVMToLLVM/NVVMToLLVM.h"
#include "mlir/include/mlir/Conversion/Passes.h"
#include "mlir/include/mlir/Conversion/UBToLLVM/UBToLLVM.h"
#include "mlir/include/mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
#include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/include/mlir/Dialect/Arith/Transforms/Passes.h"
#include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"
@ -215,6 +216,7 @@ void InitContext(mlir::MLIRContext* context) {
mlir::registerConvertMemRefToLLVMInterface(registry);
mlir::registerConvertMathToLLVMInterface(registry);
mlir::registerConvertFuncToLLVMInterface(registry);
mlir::vector::registerConvertVectorToLLVMInterface(registry);
mlir::index::registerConvertIndexToLLVMInterface(registry);
mlir::cf::registerConvertControlFlowToLLVMInterface(registry);
mlir::ub::registerConvertUBToLLVMInterface(registry);