EnzymeAD / EnzymeAD/Enzyme-JAX
[Bug] createLLVMMod crashes on parseIR due to concatenated LLVM modules from XLA
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
Summary
While building the ABI wrapper around XLA-generated LLVM IR, createLLVMMod crashes during the llvm::parseIR step with an expected top-level entity error.
The initial MHLO to LLVM compilation (compile_mhlo_to_llvm_with_xla) succeeds completely. The crash happens afterward. Because of the upstream XLA shift to Thunks, XLA is now returning multiple concatenated LLVM modules in a single string.
When createLLVMMod takes this concatenated string and passes it to llvm::parseIR to begin building the wrapper, the parser crashes as soon as it hits the second module's source_filename.
The Crash:
```
/bazel-bin/reproducer
Starting C++ createLLVMMod...
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1782238889.646514 1107501 service.cc:178] XLA service 0x14837fa0 initialized for platform Host (this does not guarantee that XLA will be used). Devices:
I0000 00:00:1782238889.646535 1107501 service.cc:194] StreamExecutor [0]: Host, Default Version (Driver: 0.0.0; Runtime: 0.0.0; Toolkit: 0.0.0; DNN: 0.0.0)
--- CRASH CONFIRMED (Inside C++) ---
failed to compile LLVM: llvmsource: :11:1: error: expected top-level entity
source_filename = "__compute_module___compute_module_wrapped_add"
^
```
Here is the llvm IR for the string returned by XLA
```llvm
; ModuleID = '__compute_module'
source_filename = "__compute_module"
target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
target triple = "x86_64-unknown-linux-gnu"
!xla_cpu_memory_region_name = !{!0}
!0 = !{!"ir_emitter"}
; ModuleID = '__compute_module___compute_module_wrapped_add'
source_filename = "__compute_module___compute_module_wrapped_add"
%XLA_CPU_KernelCallFrame = type { ptr, ptr, i64, ptr }
%XLA_CPU_KernelArg = type { ptr, i64 }
%kernel_dim3 = type { i64, i64, i64 }
define ptr @wrapped_add(ptr %0) {
%2 = getelementptr inbounds %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 3
%3 = load ptr, ptr %2, align 8, !invariant.load !2
%4 = getelementptr inbounds %XLA_CPU_KernelArg, ptr %3, i32 0, i32 0
%5 = load ptr, ptr %4, align 8, !invariant.load !2
%6 = getelementptr inbounds %XLA_CPU_KernelArg, ptr %3, i32 1, i32 0
%7 = load ptr, ptr %6, align 8, !invariant.load !2
%8 = getelementptr inbounds %XLA_CPU_KernelArg, ptr %3, i32 2, i32 0
%9 = load ptr, ptr %8, align 8, !invariant.load !2
%10 = getelementptr inbounds %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 1
%11 = load ptr, ptr %10, align 8
%12 = getelementptr inbounds %kernel_dim3, ptr %11, i32 0, i32 0
%13 = load i64, ptr %12, align 4, !invariant.load !2
%14 = call i64 @llvm.smax.i64(i64 %13, i64 0)
%15 = call i64 @llvm.smin.i64(i64 %14, i64 1)
br label %16
16: ; preds = %19, %1
%17 = phi i64 [ %20, %19 ], [ %15, %1 ]
%18 = icmp slt i64 %17, 1
br i1 %18, label %19, label %21
19: ; preds = %16
call void @wrapped_add_impl(ptr %5, ptr %5, i64 0, i64 3, i64 1, ptr %7, ptr %7, i64 0, i64 3, i64 1, ptr %9, ptr %9, i64 0, i64 3, i64 1, i64 %17)
%20 = add i64 %17, 1
br label %16
21: ; preds = %16
ret ptr null
}
; Function Attrs: alwaysinline
define internal void @wrapped_add_impl(ptr %0, ptr %1, i64 %2, i64 %3, i64 %4, ptr %5, ptr %6, i64 %7, i64 %8, i64 %9, ptr %10, ptr %11, i64 %12, i64 %13, i64 %14, i64 %15) #0 {
%17 = alloca i64, i64 4, align 8
%18 = mul i64 %3, 1
%19 = mul i64 %18, 8
%20 = getelementptr i64, ptr %1, i64 %2
call void @llvm.memcpy.p0.p0.i64(ptr %17, ptr %20, i64 %19, i1 false)
%21 = alloca i64, i64 4, align 8
%22 = mul i64 %8, 1
%23 = mul i64 %22, 8
%24 = getelementptr i64, ptr %6, i64 %7
call void @llvm.memcpy.p0.p0.i64(ptr %21, ptr %24, i64 %23, i1 false)
%25 = alloca i64, i64 4, align 64
%26 = load <4 x i64>, ptr %17, align 8
%27 = load <4 x i64>, ptr %21, align 8
%28 = add <4 x i64> %26, %27
store <4 x i64> %28, ptr %25, align 8
%29 = getelementptr i64, ptr %11, i64 %12
call void @llvm.memcpy.p0.p0.i64(ptr %29, ptr %25, i64 24, i1 false)
ret void
}
; Function Attrs: nocallback nocreateundeforpoison nofree nosync nounwind speculatable willreturn memory(none)
declare i64 @llvm.smax.i64(i64, i64) #1
; Function Attrs: nocallback nocreateundeforpoison nofree nosync nounwind speculatable willreturn memory(none)
declare i64 @llvm.smin.i64(i64, i64) #1
; Function Attrs: nocallback nofree nosync nounwind willreturn memory(argmem: readwrite)
declare void @llvm.memcpy.p0.p0.i64(ptr noalias writeonly captures(none), ptr noalias readonly captures(none), i64, i1 immarg) #2
attributes #0 = { alwaysinline }
attributes #1 = { nocallback nocreateundeforpoison nofree nosync nounwind speculatable willreturn memory(none) }
attributes #2 = { nocallback nofree nosync nounwind willreturn memory(argmem: readwrite) }
!llvm.module.flags = !{!0}
!xla_cpu_memory_region_name = !{!1}
!0 = !{i32 2, !"Debug Info Version", i32 3}
!1 = !{!"xla_cpu_emitter__tiled_emitter__hlo_opcode__fusion"}
!2 = !{}
```
We earlier used this https://github.com/EnzymeAD/Enzyme-JAX/blob/ac38bd41d87ddb90936d23cebe3d4f9fdd2bfc4f/src/enzyme_ad/jax/compile_with_xla.cc#L380
```cpp
xla::ExecutableBuildOptions build_options;
build_options.mutable_debug_options()->set_xla_embed_ir_in_executable(true);
build_options.mutable_debug_options()->set_xla_cpu_use_thunk_runtime(false);
```
Hower the set_xla_cpu_use_thunk_runtime option does not exist in XLA anymore and is now reserved in xla.proto https://github.com/openxla/xla/blob/1f93347cf23f0922b13395d5fde20881773f400d/xla/xla.proto#L1659
so we can not set that to false anymore unfortunately.
Vibecoded reproducer file:
1. `reproducer.mlir`
```mlir
module @reactant_f attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
func.func @main(%arg0: tensor<3xi64> {enzymexla.memory_effects = []}, %arg1: tensor<3xi64> {enzymexla.memory_effects = []}) -> tensor<3xi64> attributes {enzymexla.memory_effects = []} {
%0 = stablehlo.add %arg0, %arg1 : tensor<3xi64>
return %0 : tensor<3xi64>
}
}
```
2. `reproducer.cpp`
```cpp
#include
#include
#include
#include
#include
// Include the header you just provided
#include "src/enzyme_ad/jax/clang_compile.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/ArrayRef.h"
int main(int argc, char** argv) {
// 1. Load the MLIR file
std::ifstream t("reproducer.mlir");
if (!t.is_open()) {
std::cerr << "Fatal: Could not open reproducer.mlir in the current directory." << std::endl;
return 1;
}
std::stringstream buffer;
buffer << t.rdbuf();
std::string mhlo_source = buffer.str();
// 2. Set up parameters based on: %arg0: tensor<3xi64>, %arg1: tensor<3xi64> -> tensor<3xi64>
llvm::SmallVector shape_3 = {3};
// Output: One tensor of shape {3}
std::vector> out_shapes_vec = {shape_3};
llvm::ArrayRef> out_shapes(out_shapes_vec);
std::vector out_names_vec = {"out0"};
llvm::ArrayRef out_names(out_names_vec);
// Inputs: Two tensors of shape {3}
std::vector> in_shapes_vec = {shape_3, shape_3};
llvm::ArrayRef> in_shapes(in_shapes_vec);
std::vector in_names_vec = {"in0", "in1"};
llvm::ArrayRef in_names(in_names_vec);
std::vector pyargv;
std::string pass_pipeline = "";
// 3. Call the C++ function directly
std::cout << "Starting C++ createLLVMMod..." << std::endl;
auto result = createLLVMMod(
"main", // fn
mhlo_source, // source
out_shapes, // out_shapes
out_names, // out_names
in_shapes, // in_shapes
in_names, // in_names
pyargv, // pyargv_strs
ABI::Primal, // mode
Language::MHLO, // lang (Triggers the compile_mhlo_to_llvm_with_xla branch)
false, // xla_runtime
pass_pipeline // pass_pipeline
);
if (!result.ok()) {
std::cerr << "\n--- CRASH CONFIRMED (Inside C++) ---\n";
std::cerr << result.status().message() << std::endl;
return 1;
}
std::cout << "\n--- SUCCESS! ---\nIf you see this, the internal C++ function works and the bug is strictly in your extern \"C\" wrapper." << std::endl;
return 0;
}
```
3. Bazel Target:
```
cc_binary(
name = "reproducer",
srcs = ["reproducer.cpp"],
deps = [":ReactantExtraLib"],
copts = ["-std=c++20"],
)
```
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with createLLVMMod and the compile_mhlo_to_llvm_with_xla path referenced through src/enzyme_ad/jax/clang_compile.h, then run the reproducer.cpp Bazel target with reproducer.mlir. Trace how the concatenated XLA LLVM modules reach llvm::parseIR; done means the reproducer completes without the expected top-level entity error and has regression coverage for this input.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100