the specimen
12 artifacts from one run on the CPU backend · 0 levels pending · jax 0.4.38 · regenerate with python3 bench/specimen/capture.py
One program, all the way down.
one attention block, sixteen rows, no batch axis, x f32[16,32], three weights f32[32,32]. Every chapter on this site used to open a fresh toy example. This is the one program the site keeps coming back to, and what you scrub through below is not a diagram of a compiler: it is that program written down again at each level, by the machine, on a run whose command and versions are printed at the bottom of this page.
Read down the rail and watch what survives. The transpose you wrote is a real operation in StableHLO and gone by the optimized HLO, because a dot contracts along whichever dimension it is told to. The square root of thirty-two is an instruction at one level and the literal 0.176776692 at the next. Five dot products go into the backend and none of them come out as LLVM IR, since the CPU backend calls a library kernel instead of generating one.
explorer
the specimen, level by level 01 / 12
"""The specimen: one attention block, and the exact arrays it runs on.
Sixteen rows, thirty-two channels, no batch axis. Three projections, one
scaled dot product, one softmax, one more dot product. That is the smallest
program that still contains everything the site teaches at the compiler
floor: matmuls that get their operands transposed for free, a constant that
folds, and a softmax that arrives as nine primitives instead of one.
Nothing else in this directory decides anything about the program. The
capture script imports this module and writes down what the machine makes of
it, level by level.
"""
import jax
import jax.numpy as jnp
ROWS = 16
CHANNELS = 32
SEED = 0
def block(x, wq, wk, wv):
q, k, v = x @ wq, x @ wk, x @ wv
s = (q @ k.T) / jnp.sqrt(jnp.float32(q.shape[-1]))
return jax.nn.softmax(s, axis=-1) @ v
def inputs():
"""The four arrays the specimen is captured on, from one fixed seed.
Values do not change any artifact below the source: the trace and every
level under it depend on shapes and dtypes only. They are seeded anyway
so that running the block for its numbers gives the same numbers twice.
"""
kx, kq, kk, kv = jax.random.split(jax.random.PRNGKey(SEED), 4)
x = jax.random.normal(kx, (ROWS, CHANNELS), jnp.float32)
wq = jax.random.normal(kq, (CHANNELS, CHANNELS), jnp.float32)
wk = jax.random.normal(kk, (CHANNELS, CHANNELS), jnp.float32)
wv = jax.random.normal(kv, (CHANNELS, CHANNELS), jnp.float32)
return x, wq, wk, wvbench/specimen/artifacts/source.py
{ lambda ; a:f32[16,32] b:f32[32,32] c:f32[32,32] d:f32[32,32]. let
e:f32[16,32] = dot_general[
dimension_numbers=(([1], [0]), ([], []))
preferred_element_type=float32
] a b
f:f32[16,32] = dot_general[
dimension_numbers=(([1], [0]), ([], []))
preferred_element_type=float32
] a c
g:f32[16,32] = dot_general[
dimension_numbers=(([1], [0]), ([], []))
preferred_element_type=float32
] a d
h:f32[32,16] = transpose[permutation=(1, 0)] f
i:f32[16,16] = dot_general[
dimension_numbers=(([1], [0]), ([], []))
preferred_element_type=float32
] e h
j:f32[] = sqrt 32.0
k:f32[16,16] = div i j
l:f32[16] = reduce_max[axes=(1,)] k
m:f32[16] = max -inf l
n:f32[16,1] = broadcast_in_dim[
broadcast_dimensions=(0,)
shape=(16, 1)
sharding=None
] m
o:f32[16,1] = stop_gradient n
p:f32[16,16] = sub k o
q:f32[16,16] = exp p
r:f32[16] = reduce_sum[axes=(1,)] q
s:f32[16,1] = broadcast_in_dim[
broadcast_dimensions=(0,)
shape=(16, 1)
sharding=None
] r
t:f32[16,16] = div q s
u:f32[16,32] = dot_general[
dimension_numbers=(([1], [0]), ([], []))
preferred_element_type=float32
] t g
in (u,) }bench/specimen/artifacts/jaxpr.txt
== lazy IR (torch_xla._XLAC._get_xla_tensors_text) ==
IR {
%0 = f32[32,32]{1,0} xla::device_data(), xla_shape=f32[32,32]{1,0}
%1 = f32[16,32]{1,0} xla::device_data(), xla_shape=f32[16,32]{1,0}
%2 = f32[16,32]{1,0} aten::mm(%1, %0), xla_shape=f32[16,32]{1,0}
%3 = f32[] xla::device_data(), xla_shape=f32[]
%4 = f32[] aten::sqrt(%3), xla_shape=f32[]
%5 = f32[32,32]{1,0} xla::device_data(), xla_shape=f32[32,32]{1,0}
%6 = f32[16,32]{1,0} aten::mm(%1, %5), xla_shape=f32[16,32]{1,0}
%7 = f32[32,16]{0,1} aten::permute(%6), xla_shape=f32[32,16]{0,1}
%8 = f32[32,32]{1,0} xla::device_data(), xla_shape=f32[32,32]{1,0}
%9 = f32[16,32]{1,0} aten::mm(%1, %8), xla_shape=f32[16,32]{1,0}
%10 = f32[16,16]{1,0} aten::mm(%9, %7), xla_shape=f32[16,16]{1,0}
%11 = f32[16,16]{1,0} aten::div(%10, %4), xla_shape=f32[16,16]{1,0}
%12 = f32[16,16]{1,0} aten::softmax(%11), xla_shape=f32[16,16]{1,0}
%13 = f32[16,32]{1,0} aten::mm(%12, %2), xla_shape=f32[16,32]{1,0}, ROOT=0
}
== HLO (torch_xla._XLAC._get_xla_tensors_hlo) ==
HloModule IrToHlo.33, entry_computation_layout={(f32[32,32]{1,0}, f32[16,32]{1,0}, f32[], f32[32,32]{1,0}, f32[32,32]{1,0})->(f32[16,32]{1,0})}
%MaxComputation.15 (x.16: f32[], y.17: f32[]) -> f32[] {
%x.16 = f32[] parameter(0)
%y.17 = f32[] parameter(1)
ROOT %maximum.18 = f32[] maximum(f32[] %x.16, f32[] %y.17)
}
%AddComputation.24 (x.25: f32[], y.26: f32[]) -> f32[] {
%x.25 = f32[] parameter(0)
%y.26 = f32[] parameter(1)
ROOT %add.27 = f32[] add(f32[] %x.25, f32[] %y.26)
}
ENTRY %IrToHlo.33 (p0.1: f32[32,32], p1.2: f32[16,32], p2.4: f32[], p3.6: f32[32,32], p4.9: f32[32,32]) -> (f32[16,32]) {
%p1.2 = f32[16,32]{1,0} parameter(1)
%p4.9 = f32[32,32]{1,0} parameter(4)
%dot.10 = f32[16,32]{1,0} dot(f32[16,32]{1,0} %p1.2, f32[32,32]{1,0} %p4.9), lhs_contracting_dims={1}, rhs_contracting_dims={0}
%p3.6 = f32[32,32]{1,0} parameter(3)
%dot.7 = f32[16,32]{1,0} dot(f32[16,32]{1,0} %p1.2, f32[32,32]{1,0} %p3.6), lhs_contracting_dims={1}, rhs_contracting_dims={0}
%transpose.8 = f32[32,16]{0,1} transpose(f32[16,32]{1,0} %dot.7), dimensions={1,0}
%dot.11 = f32[16,16]{1,0} dot(f32[16,32]{1,0} %dot.10, f32[32,16]{0,1} %transpose.8), lhs_contracting_dims={1}, rhs_contracting_dims={0}
%p2.4 = f32[] parameter(2)
%sqrt.5 = f32[] sqrt(f32[] %p2.4)
%broadcast.12 = f32[16,16]{1,0} broadcast(f32[] %sqrt.5), dimensions={}
%divide.13 = f32[16,16]{1,0} divide(f32[16,16]{1,0} %dot.11, f32[16,16]{1,0} %broadcast.12)
%constant.14 = f32[] constant(-inf)
%reduce.19 = f32[16]{0} reduce(f32[16,16]{1,0} %divide.13, f32[] %constant.14), dimensions={1}, to_apply=%MaxComputation.15
%broadcast.20 = f32[16,16]{1,0} broadcast(f32[16]{0} %reduce.19), dimensions={0}
%subtract.21 = f32[16,16]{1,0} subtract(f32[16,16]{1,0} %divide.13, f32[16,16]{1,0} %broadcast.20)
%exponential.22 = f32[16,16]{1,0} exponential(f32[16,16]{1,0} %subtract.21)
%constant.23 = f32[] constant(0)
%reduce.28 = f32[16]{0} reduce(f32[16,16]{1,0} %exponential.22, f32[] %constant.23), dimensions={1}, to_apply=%AddComputation.24
%broadcast.29 = f32[16,16]{1,0} broadcast(f32[16]{0} %reduce.28), dimensions={0}
%divide.30 = f32[16,16]{1,0} divide(f32[16,16]{1,0} %exponential.22, f32[16,16]{1,0} %broadcast.29)
%p0.1 = f32[32,32]{1,0} parameter(0)
%dot.3 = f32[16,32]{1,0} dot(f32[16,32]{1,0} %p1.2, f32[32,32]{1,0} %p0.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
%dot.31 = f32[16,32]{1,0} dot(f32[16,16]{1,0} %divide.30, f32[16,32]{1,0} %dot.3), lhs_contracting_dims={1}, rhs_contracting_dims={0}
ROOT %tuple.32 = (f32[16,32]{1,0}) tuple(f32[16,32]{1,0} %dot.31)
}
== seam (torch_xla.debug.metrics, TotalSamples after one sync and one read) ==
CompileTime 1
ExecuteTime 2
TransferToDeviceTime 1
TransferFromDeviceTime 1
OutboundData 1
InboundData 1bench/specimen/artifacts/torch-trace.txt
module @jit_block attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
func.func public @main(%arg0: tensor<16x32xf32>, %arg1: tensor<32x32xf32>, %arg2: tensor<32x32xf32>, %arg3: tensor<32x32xf32>) -> (tensor<16x32xf32> {jax.result_info = ""}) {
%0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<16x32xf32>, tensor<32x32xf32>) -> tensor<16x32xf32>
%1 = stablehlo.dot_general %arg0, %arg2, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<16x32xf32>, tensor<32x32xf32>) -> tensor<16x32xf32>
%2 = stablehlo.dot_general %arg0, %arg3, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<16x32xf32>, tensor<32x32xf32>) -> tensor<16x32xf32>
%3 = stablehlo.transpose %1, dims = [1, 0] : (tensor<16x32xf32>) -> tensor<32x16xf32>
%4 = stablehlo.dot_general %0, %3, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<16x32xf32>, tensor<32x16xf32>) -> tensor<16x16xf32>
%cst = stablehlo.constant dense<3.200000e+01> : tensor<f32>
%5 = stablehlo.sqrt %cst : tensor<f32>
%6 = stablehlo.broadcast_in_dim %5, dims = [] : (tensor<f32>) -> tensor<16x16xf32>
%7 = stablehlo.divide %4, %6 : tensor<16x16xf32>
%cst_0 = stablehlo.constant dense<0xFF800000> : tensor<f32>
%8 = stablehlo.reduce(%7 init: %cst_0) applies stablehlo.maximum across dimensions = [1] : (tensor<16x16xf32>, tensor<f32>) -> tensor<16xf32>
%cst_1 = stablehlo.constant dense<0xFF800000> : tensor<f32>
%9 = stablehlo.broadcast_in_dim %cst_1, dims = [] : (tensor<f32>) -> tensor<16xf32>
%10 = stablehlo.maximum %9, %8 : tensor<16xf32>
%11 = stablehlo.broadcast_in_dim %10, dims = [0] : (tensor<16xf32>) -> tensor<16x1xf32>
%12 = stablehlo.broadcast_in_dim %11, dims = [0, 1] : (tensor<16x1xf32>) -> tensor<16x16xf32>
%13 = stablehlo.subtract %7, %12 : tensor<16x16xf32>
%14 = stablehlo.exponential %13 : tensor<16x16xf32>
%cst_2 = stablehlo.constant dense<0.000000e+00> : tensor<f32>
%15 = stablehlo.reduce(%14 init: %cst_2) applies stablehlo.add across dimensions = [1] : (tensor<16x16xf32>, tensor<f32>) -> tensor<16xf32>
%16 = stablehlo.broadcast_in_dim %15, dims = [0] : (tensor<16xf32>) -> tensor<16x1xf32>
%17 = stablehlo.broadcast_in_dim %16, dims = [0, 1] : (tensor<16x1xf32>) -> tensor<16x16xf32>
%18 = stablehlo.divide %14, %17 : tensor<16x16xf32>
%19 = stablehlo.dot_general %18, %2, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<16x16xf32>, tensor<16x32xf32>) -> tensor<16x32xf32>
return %19 : tensor<16x32xf32>
}
}
bench/specimen/artifacts/stablehlo.mlir
HloModule jit_block, entry_computation_layout={(f32[16,32]{1,0}, f32[32,32]{1,0}, f32[32,32]{1,0}, f32[32,32]{1,0})->f32[16,32]{1,0}}, allow_spmd_sharding_propagation_to_parameters={true,true,true,true}, allow_spmd_sharding_propagation_to_output={true}
region_0.17 {
Arg_0.18 = f32[] parameter(0), metadata={op_name="jit(block)/jit(main)/reduce_max"}
Arg_1.19 = f32[] parameter(1), metadata={op_name="jit(block)/jit(main)/reduce_max"}
ROOT maximum.20 = f32[] maximum(Arg_0.18, Arg_1.19), metadata={op_name="jit(block)/jit(main)/reduce_max"}
}
region_1.29 {
Arg_0.30 = f32[] parameter(0), metadata={op_name="jit(block)/jit(main)/reduce_sum"}
Arg_1.31 = f32[] parameter(1), metadata={op_name="jit(block)/jit(main)/reduce_sum"}
ROOT add.32 = f32[] add(Arg_0.30, Arg_1.31), metadata={op_name="jit(block)/jit(main)/reduce_sum"}
}
ENTRY main.40 {
Arg_0.1 = f32[16,32]{1,0} parameter(0), metadata={op_name="x"}
Arg_1.2 = f32[32,32]{1,0} parameter(1), metadata={op_name="wq"}
dot.11 = f32[16,32]{1,0} dot(Arg_0.1, Arg_1.2), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(block)/jit(main)/dot_general"}
Arg_2.3 = f32[32,32]{1,0} parameter(2), metadata={op_name="wk"}
dot.12 = f32[16,32]{1,0} dot(Arg_0.1, Arg_2.3), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(block)/jit(main)/dot_general"}
transpose.14 = f32[32,16]{0,1} transpose(dot.12), dimensions={1,0}, metadata={op_name="jit(block)/jit(main)/transpose"}
dot.15 = f32[16,16]{1,0} dot(dot.11, transpose.14), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(block)/jit(main)/dot_general"}
constant.7 = f32[] constant(5.65685415)
broadcast.8 = f32[16,16]{1,0} broadcast(constant.7), dimensions={}
divide.16 = f32[16,16]{1,0} divide(dot.15, broadcast.8), metadata={op_name="jit(block)/jit(main)/div"}
constant.10 = f32[] constant(-inf)
reduce.21 = f32[16]{0} reduce(divide.16, constant.10), dimensions={1}, to_apply=region_0.17, metadata={op_name="jit(block)/jit(main)/reduce_max"}
constant.5 = f32[] constant(-inf)
broadcast.6 = f32[16]{0} broadcast(constant.5), dimensions={}
maximum.22 = f32[16]{0} maximum(reduce.21, broadcast.6), metadata={op_name="jit(block)/jit(main)/max"}
reshape.23 = f32[16,1]{1,0} reshape(maximum.22), metadata={op_name="jit(block)/jit(main)/broadcast_in_dim"}
broadcast.24 = f32[16,1]{1,0} broadcast(reshape.23), dimensions={0,1}, metadata={op_name="jit(block)/jit(main)/sub"}
reshape.25 = f32[16]{0} reshape(broadcast.24), metadata={op_name="jit(block)/jit(main)/sub"}
broadcast.26 = f32[16,16]{1,0} broadcast(reshape.25), dimensions={0}, metadata={op_name="jit(block)/jit(main)/sub"}
subtract.27 = f32[16,16]{1,0} subtract(divide.16, broadcast.26), metadata={op_name="jit(block)/jit(main)/sub"}
exponential.28 = f32[16,16]{1,0} exponential(subtract.27), metadata={op_name="jit(block)/jit(main)/exp"}
constant.9 = f32[] constant(0)
reduce.33 = f32[16]{0} reduce(exponential.28, constant.9), dimensions={1}, to_apply=region_1.29, metadata={op_name="jit(block)/jit(main)/reduce_sum"}
reshape.34 = f32[16,1]{1,0} reshape(reduce.33), metadata={op_name="jit(block)/jit(main)/broadcast_in_dim"}
broadcast.35 = f32[16,1]{1,0} broadcast(reshape.34), dimensions={0,1}, metadata={op_name="jit(block)/jit(main)/div"}
reshape.36 = f32[16]{0} reshape(broadcast.35), metadata={op_name="jit(block)/jit(main)/div"}
broadcast.37 = f32[16,16]{1,0} broadcast(reshape.36), dimensions={0}, metadata={op_name="jit(block)/jit(main)/div"}
divide.38 = f32[16,16]{1,0} divide(exponential.28, broadcast.37), metadata={op_name="jit(block)/jit(main)/div"}
Arg_3.4 = f32[32,32]{1,0} parameter(3), metadata={op_name="wv"}
dot.13 = f32[16,32]{1,0} dot(Arg_0.1, Arg_3.4), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(block)/jit(main)/dot_general"}
ROOT dot.39 = f32[16,32]{1,0} dot(divide.38, dot.13), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(block)/jit(main)/dot_general"}
} // main.40
bench/specimen/artifacts/hlo-before-optimizations.txt
HloModule jit_block, is_scheduled=true, entry_computation_layout={(f32[16,32]{1,0}, f32[32,32]{1,0}, f32[32,32]{1,0}, f32[32,32]{1,0})->f32[16,32]{1,0}}, allow_spmd_sharding_propagation_to_parameters={true,true,true,true}, allow_spmd_sharding_propagation_to_output={true}
region_1.29 {
Arg_0.30 = f32[] parameter(0), metadata={op_name="jit(block)/jit(main)/reduce_sum"}
Arg_1.31 = f32[] parameter(1), metadata={op_name="jit(block)/jit(main)/reduce_sum"}
ROOT add.32 = f32[] add(Arg_0.30, Arg_1.31), metadata={op_name="jit(block)/jit(main)/reduce_sum"}
}
fused_computation {
param_0 = f32[16,16]{1,0} parameter(0)
param_1.1 = f32[16]{0} parameter(1)
broadcast.3 = f32[16,16]{1,0} broadcast(param_1.1), dimensions={0}, metadata={op_name="jit(block)/jit(main)/div"}
ROOT divide.0 = f32[16,16]{1,0} divide(param_0, broadcast.3), metadata={op_name="jit(block)/jit(main)/div"}
}
fused_computation.1 {
param_1.5 = f32[16,16]{1,0} parameter(1)
constant.2 = f32[] constant(0.176776692)
broadcast.7 = f32[16,16]{1,0} broadcast(constant.2), dimensions={}
multiply.1 = f32[16,16]{1,0} multiply(param_1.5, broadcast.7), metadata={op_name="jit(block)/jit(main)/div"}
param_0.3 = f32[16]{0} parameter(0)
constant.1 = f32[] constant(-inf)
broadcast.5 = f32[16]{0} broadcast(constant.1), dimensions={}
maximum.0 = f32[16]{0} maximum(param_0.3, broadcast.5), metadata={op_name="jit(block)/jit(main)/max"}
broadcast.4 = f32[16,16]{1,0} broadcast(maximum.0), dimensions={0}, metadata={op_name="jit(block)/jit(main)/sub"}
subtract.0 = f32[16,16]{1,0} subtract(multiply.1, broadcast.4), metadata={op_name="jit(block)/jit(main)/sub"}
ROOT exponential.0 = f32[16,16]{1,0} exponential(subtract.0), metadata={op_name="jit(block)/jit(main)/exp"}
} // fused_computation.1
region_0.17 {
Arg_0.18 = f32[] parameter(0), metadata={op_name="jit(block)/jit(main)/reduce_max"}
Arg_1.19 = f32[] parameter(1), metadata={op_name="jit(block)/jit(main)/reduce_max"}
ROOT maximum.20 = f32[] maximum(Arg_0.18, Arg_1.19), metadata={op_name="jit(block)/jit(main)/reduce_max"}
}
fused_computation.2 {
param_0.6 = f32[16,16]{1,0} parameter(0)
constant.4 = f32[] constant(0.176776692)
broadcast.9 = f32[16,16]{1,0} broadcast(constant.4), dimensions={}
multiply.2 = f32[16,16]{1,0} multiply(param_0.6, broadcast.9), metadata={op_name="jit(block)/jit(main)/div"}
constant.3 = f32[] constant(-inf)
ROOT reduce.0 = f32[16]{0} reduce(multiply.2, constant.3), dimensions={1}, to_apply=region_0.17, metadata={op_name="jit(block)/jit(main)/reduce_max"}
} // fused_computation.2
ENTRY main.40 {
Arg_0.1 = f32[16,32]{1,0} parameter(0), metadata={op_name="x"}
Arg_2.3 = f32[32,32]{1,0} parameter(2), metadata={op_name="wk"}
dot = f32[32,16]{1,0} dot(Arg_2.3, Arg_0.1), lhs_contracting_dims={0}, rhs_contracting_dims={1}, metadata={op_name="jit(block)/jit(main)/transpose"}
Arg_1.2 = f32[32,32]{1,0} parameter(1), metadata={op_name="wq"}
dot.11 = f32[16,32]{1,0} dot(Arg_0.1, Arg_1.2), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(block)/jit(main)/dot_general"}
dot.15 = f32[16,16]{1,0} dot(dot.11, dot), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(block)/jit(main)/dot_general"}
multiply_reduce_fusion = f32[16]{0} fusion(dot.15), kind=kLoop, calls=fused_computation.2, metadata={op_name="jit(block)/jit(main)/reduce_max"}
subtract_exponential_fusion = f32[16,16]{1,0} fusion(multiply_reduce_fusion, dot.15), kind=kLoop, calls=fused_computation.1, metadata={op_name="jit(block)/jit(main)/exp"}
constant.9 = f32[] constant(0)
reduce.33 = f32[16]{0} reduce(subtract_exponential_fusion, constant.9), dimensions={1}, to_apply=region_1.29, metadata={op_name="jit(block)/jit(main)/reduce_sum"}
broadcast_divide_fusion = f32[16,16]{1,0} fusion(subtract_exponential_fusion, reduce.33), kind=kLoop, calls=fused_computation, metadata={op_name="jit(block)/jit(main)/div"}
Arg_3.4 = f32[32,32]{1,0} parameter(3), metadata={op_name="wv"}
dot.13 = f32[16,32]{1,0} dot(Arg_0.1, Arg_3.4), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(block)/jit(main)/dot_general"}
ROOT dot.39 = f32[16,32]{1,0} dot(broadcast_divide_fusion, dot.13), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(block)/jit(main)/dot_general"}
} // main.40
bench/specimen/artifacts/hlo-after-optimizations.txt
module {
tt.func @matmul_kernel__Pfp32_Pfp32_Pfp32_i32_i32_i32_i32_i32_i32_i32_i32_i32__12c64_13c64_14c64_15c8(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32 {tt.divisibility = 16 : i32}, %arg7: i32, %arg8: i32 {tt.divisibility = 16 : i32}, %arg9: i32, %arg10: i32 {tt.divisibility = 16 : i32}, %arg11: i32) {
%cst = arith.constant dense<true> : tensor<64x64xi1>
%c64 = arith.constant 64 : i32
%c0 = arith.constant 0 : i32
%cst_0 = arith.constant dense<0.000000e+00> : tensor<64x64xf32>
%c64_i32 = arith.constant 64 : i32
%c63_i32 = arith.constant 63 : i32
%c8_i32 = arith.constant 8 : i32
%0 = tt.get_program_id x : i32
%1 = arith.addi %arg3, %c63_i32 : i32
%2 = arith.divsi %1, %c64_i32 : i32
%3 = arith.addi %arg4, %c63_i32 : i32
%4 = arith.divsi %3, %c64_i32 : i32
%5 = arith.muli %4, %c8_i32 : i32
%6 = arith.divsi %0, %5 : i32
%7 = arith.muli %6, %c8_i32 : i32
%8 = arith.subi %2, %7 : i32
%9 = arith.cmpi slt, %8, %c8_i32 : i32
%10 = arith.select %9, %8, %c8_i32 : i32
%11 = arith.remsi %0, %10 : i32
%12 = arith.addi %7, %11 : i32
%13 = arith.remsi %0, %5 : i32
%14 = arith.divsi %13, %10 : i32
%15 = arith.muli %12, %c64_i32 : i32
%16 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
%17 = tt.splat %15 : i32 -> tensor<64xi32>
%18 = arith.addi %17, %16 : tensor<64xi32>
%19 = arith.muli %14, %c64_i32 : i32
%20 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
%21 = tt.splat %19 : i32 -> tensor<64xi32>
%22 = arith.addi %21, %20 : tensor<64xi32>
%23 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32>
%24 = tt.expand_dims %18 {axis = 1 : i32} : tensor<64xi32> -> tensor<64x1xi32>
%25 = tt.splat %arg6 : i32 -> tensor<64x1xi32>
%26 = arith.muli %24, %25 : tensor<64x1xi32>
%27 = tt.expand_dims %23 {axis = 0 : i32} : tensor<64xi32> -> tensor<1x64xi32>
%28 = tt.splat %arg7 : i32 -> tensor<1x64xi32>
%29 = arith.muli %27, %28 : tensor<1x64xi32>
%30 = tt.broadcast %26 : tensor<64x1xi32> -> tensor<64x64xi32>
%31 = tt.broadcast %29 : tensor<1x64xi32> -> tensor<64x64xi32>
%32 = arith.addi %30, %31 : tensor<64x64xi32>
%33 = tt.splat %arg0 : !tt.ptr<f32> -> tensor<64x64x!tt.ptr<f32>>
%34 = tt.addptr %33, %32 : tensor<64x64x!tt.ptr<f32>>, tensor<64x64xi32>
%35 = tt.expand_dims %23 {axis = 1 : i32} : tensor<64xi32> -> tensor<64x1xi32>
%36 = tt.splat %arg8 : i32 -> tensor<64x1xi32>
%37 = arith.muli %35, %36 : tensor<64x1xi32>
%38 = tt.expand_dims %22 {axis = 0 : i32} : tensor<64xi32> -> tensor<1x64xi32>
%39 = tt.splat %arg9 : i32 -> tensor<1x64xi32>
%40 = arith.muli %38, %39 : tensor<1x64xi32>
%41 = tt.broadcast %37 : tensor<64x1xi32> -> tensor<64x64xi32>
%42 = tt.broadcast %40 : tensor<1x64xi32> -> tensor<64x64xi32>
%43 = arith.addi %41, %42 : tensor<64x64xi32>
%44 = tt.splat %arg1 : !tt.ptr<f32> -> tensor<64x64x!tt.ptr<f32>>
%45 = tt.addptr %44, %43 : tensor<64x64x!tt.ptr<f32>>, tensor<64x64xi32>
%47:3 = scf.for %arg12 = %c0 to %arg5 step %c64 iter_args(%arg13 = %cst_0, %arg14 = %34, %arg15 = %45) -> (tensor<64x64xf32>, tensor<64x64x!tt.ptr<f32>>, tensor<64x64x!tt.ptr<f32>>) : i32 {
%76 = tt.load %arg14, %cst, %cst_0 : tensor<64x64x!tt.ptr<f32>>
%77 = tt.load %arg15, %cst, %cst_0 : tensor<64x64x!tt.ptr<f32>>
%78 = tt.dot %76, %77, %cst_0 : tensor<64x64xf32> * tensor<64x64xf32> -> tensor<64x64xf32>
%79 = arith.addf %arg13, %78 : tensor<64x64xf32>
%80 = arith.muli %arg7, %c64_i32 : i32
%81 = tt.splat %80 : i32 -> tensor<64x64xi32>
%82 = tt.addptr %arg14, %81 : tensor<64x64x!tt.ptr<f32>>, tensor<64x64xi32>
%83 = arith.muli %arg8, %c64_i32 : i32
%84 = tt.splat %83 : i32 -> tensor<64x64xi32>
%85 = tt.addptr %arg15, %84 : tensor<64x64x!tt.ptr<f32>>, tensor<64x64xi32>
scf.yield %79, %82, %85 : tensor<64x64xf32>, tensor<64x64x!tt.ptr<f32>>, tensor<64x64x!tt.ptr<f32>>
}
bench/specimen/artifacts/triton-ttir.mlir
#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [32, 1], warpsPerCTA = [2, 2], order = [0, 1]}>
#blocked1 = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [32, 1], warpsPerCTA = [1, 4], order = [0, 1]}>
#mma = #ttg.nvidia_mma<{versionMajor = 2, versionMinor = 0, warpsPerCTA = [4, 1], instrShape = [16, 8]}>
#shared = #ttg.swizzled_shared<{vec = 8, perPhase = 1, maxPhase = 4, order = [0, 1]}>
#shared1 = #ttg.swizzled_shared<{vec = 8, perPhase = 1, maxPhase = 4, order = [1, 0]}>
#smem = #ttg.shared_memory
module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
tt.func public @_jagged_hstu_attn_fwd_0d1d2d3d4d5de(%arg0: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg3: !tt.ptr<i64> {tt.divisibility = 16 : i32}, %arg4: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg5: i32 {tt.divisibility = 16 : i32, tt.max_divisibility = 8 : i32}) {
%cst = arith.constant dense<0.000000e+00> : tensor<64x32xf32, #mma>
%c64_i32 = arith.constant 64 : i32
%c0_i32 = arith.constant 0 : i32
%c32_i32 = arith.constant 32 : i32
%0 = tt.get_program_id x : i32
%1 = arith.muli %0, %c64_i32 : i32
%2 = tt.get_program_id y : i32
%3 = tt.load %arg3 : !tt.ptr<i64>
%4 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #ttg.slice<{dim = 1, parent = #blocked}>>
%5 = tt.splat %1 : i32 -> tensor<64xi32, #ttg.slice<{dim = 1, parent = #blocked}>>
%6 = arith.addi %5, %4 : tensor<64xi32, #ttg.slice<{dim = 1, parent = #blocked}>>
%7 = tt.expand_dims %6 {axis = 1 : i32} : tensor<64xi32, #ttg.slice<{dim = 1, parent = #blocked}>> -> tensor<64x1xi32, #blocked>
%8 = tt.splat %3 : i64 -> tensor<64x1xi64, #blocked>
%9 = arith.extsi %7 : tensor<64x1xi32, #blocked> to tensor<64x1xi64, #blocked>
%10 = arith.addi %8, %9 : tensor<64x1xi64, #blocked>
%11 = arith.extsi %arg5 : i32 to i64
%12 = tt.splat %11 : i64 -> tensor<64x1xi64, #blocked>
%13 = arith.muli %10, %12 : tensor<64x1xi64, #blocked>
%14 = arith.muli %2, %arg5 : i32
%15 = arith.extsi %14 : i32 to i64
%16 = tt.splat %15 : i64 -> tensor<64x1xi64, #blocked>
%17 = arith.addi %13, %16 : tensor<64x1xi64, #blocked>
%18 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #ttg.slice<{dim = 0, parent = #blocked}>>
%19 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #ttg.slice<{dim = 0, parent = #blocked1}>>
%20 = tt.expand_dims %18 {axis = 0 : i32} : tensor<64xi32, #ttg.slice<{dim = 0, parent = #blocked}>> -> tensor<1x64xi32, #blocked>
%21 = tt.expand_dims %19 {axis = 0 : i32} : tensor<64xi32, #ttg.slice<{dim = 0, parent = #blocked1}>> -> tensor<1x64xi32, #blocked1>
%22 = tt.splat %arg5 : i32 -> tensor<1x64xi32, #blocked>
%23 = tt.splat %arg5 : i32 -> tensor<1x64xi32, #blocked1>
%24 = arith.muli %20, %22 : tensor<1x64xi32, #blocked>
%25 = arith.muli %21, %23 : tensor<1x64xi32, #blocked1>
%26 = tt.broadcast %17 : tensor<64x1xi64, #blocked> -> tensor<64x64xi64, #blocked>
%27 = arith.extsi %24 : tensor<1x64xi32, #blocked> to tensor<1x64xi64, #blocked>
%28 = arith.extsi %25 : tensor<1x64xi32, #blocked1> to tensor<1x64xi64, #blocked1>
%29 = tt.broadcast %27 : tensor<1x64xi64, #blocked> -> tensor<64x64xi64, #blocked>
%30 = arith.addi %26, %29 : tensor<64x64xi64, #blocked>
%31 = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<32xi32, #ttg.slice<{dim = 1, parent = #blocked1}>>
%32 = tt.expand_dims %31 {axis = 1 : i32} : tensor<32xi32, #ttg.slice<{dim = 1, parent = #blocked1}>> -> tensor<32x1xi32, #blocked1>
%33 = tt.splat %3 : i64 -> tensor<32x1xi64, #blocked1>
%34 = arith.extsi %32 : tensor<32x1xi32, #blocked1> to tensor<32x1xi64, #blocked1>
%35 = arith.addi %33, %34 : tensor<32x1xi64, #blocked1>
%36 = tt.splat %11 : i64 -> tensor<32x1xi64, #blocked1>
%37 = arith.muli %35, %36 : tensor<32x1xi64, #blocked1>
%38 = tt.splat %15 : i64 -> tensor<32x1xi64, #blocked1>
%39 = arith.addi %37, %38 : tensor<32x1xi64, #blocked1>
%40 = tt.broadcast %39 : tensor<32x1xi64, #blocked1> -> tensor<32x64xi64, #blocked1>
%41 = tt.broadcast %28 : tensor<1x64xi64, #blocked1> -> tensor<32x64xi64, #blocked1>
%42 = arith.addi %40, %41 : tensor<32x64xi64, #blocked1>
%43 = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<32xi32, #ttg.slice<{dim = 0, parent = #blocked1}>>
%44 = tt.make_range {end = 32 : i32, start = 0 : i32} : tensor<32xi32, #ttg.slice<{dim = 0, parent = #blocked}>>
%45 = tt.expand_dims %43 {axis = 0 : i32} : tensor<32xi32, #ttg.slice<{dim = 0, parent = #blocked1}>> -> tensor<1x32xi32, #blocked1>
%46 = tt.expand_dims %44 {axis = 0 : i32} : tensor<32xi32, #ttg.slice<{dim = 0, parent = #blocked}>> -> tensor<1x32xi32, #blocked>
%47 = tt.splat %arg5 : i32 -> tensor<1x32xi32, #blocked1>
%48 = tt.splat %arg5 : i32 -> tensor<1x32xi32, #blocked>
%49 = arith.muli %45, %47 : tensor<1x32xi32, #blocked1>
%50 = arith.muli %46, %48 : tensor<1x32xi32, #blocked>
%51 = tt.broadcast %39 : tensor<32x1xi64, #blocked1> -> tensor<32x32xi64, #blocked1>
%52 = arith.extsi %49 : tensor<1x32xi32, #blocked1> to tensor<1x32xi64, #blocked1>
%53 = arith.extsi %50 : tensor<1x32xi32, #blocked> to tensor<1x32xi64, #blocked>
%54 = tt.broadcast %52 : tensor<1x32xi64, #blocked1> -> tensor<32x32xi64, #blocked1>
%55 = arith.addi %51, %54 : tensor<32x32xi64, #blocked1>
%56 = tt.splat %arg0 : !tt.ptr<f32> -> tensor<64x64x!tt.ptr<f32>, #blocked>
%57 = tt.addptr %56, %30 : tensor<64x64x!tt.ptr<f32>, #blocked>, tensor<64x64xi64, #blocked>
%58 = tt.splat %arg1 : !tt.ptr<f32> -> tensor<32x64x!tt.ptr<f32>, #blocked1>
%59 = tt.addptr %58, %42 : tensor<32x64x!tt.ptr<f32>, #blocked1>, tensor<32x64xi64, #blocked1>
%60 = tt.splat %arg2 : !tt.ptr<f32> -> tensor<32x32x!tt.ptr<f32>, #blocked1>
%61 = tt.addptr %60, %55 : tensor<32x32x!tt.ptr<f32>, #blocked1>, tensor<32x32xi64, #blocked1>
%62 = tt.load %57 : tensor<64x64x!tt.ptr<f32>, #blocked>
%63 = scf.for %arg6 = %c0_i32 to %c64_i32 step %c32_i32 iter_args(%arg7 = %cst) -> (tensor<64x32xf32, #mma>) : i32 {
%70 = tt.load %59 : tensor<32x64x!tt.ptr<f32>, #blocked1>
%71 = ttg.convert_layout %62 : tensor<64x64xf32, #blocked> -> tensor<64x64xf32, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 1}>>
%72 = ttg.local_alloc %70 : (tensor<32x64xf32, #blocked1>) -> !ttg.memdesc<32x64xf32, #shared, #smem>
%73 = ttg.memdesc_trans %72 {order=array<i32: 1,0>} : !ttg.memdesc<32x64xf32, #shared, #smem> -> !ttg.memdesc<64x32xf32, #shared1, #smem>
%74 = ttg.local_load %73 : !ttg.memdesc<64x32xf32, #shared1, #smem> -> tensor<64x32xf32, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 1}>>
%75 = tt.dot %71, %74, %cst, inputPrecision = tf32 : tensor<64x64xf32, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 1}>> * tensor<64x32xf32, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 1}>> -> tensor<64x32xf32, #mma>
%76 = tt.load %61 : tensor<32x32x!tt.ptr<f32>, #blocked1>
%77 = ttg.convert_layout %75 : tensor<64x32xf32, #mma> -> tensor<64x32xf32, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 1}>>
%78 = ttg.convert_layout %76 : tensor<32x32xf32, #blocked1> -> tensor<32x32xf32, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 1}>>
%79 = tt.dot %77, %78, %arg7, inputPrecision = tf32 : tensor<64x32xf32, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 1}>> * tensor<32x32xf32, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 1}>> -> tensor<64x32xf32, #mma>
scf.yield %79 : tensor<64x32xf32, #mma>
}
%64 = tt.broadcast %17 : tensor<64x1xi64, #blocked> -> tensor<64x32xi64, #blocked>
%65 = tt.broadcast %53 : tensor<1x32xi64, #blocked> -> tensor<64x32xi64, #blocked>
%66 = arith.addi %64, %65 : tensor<64x32xi64, #blocked>
%67 = tt.splat %arg4 : !tt.ptr<f32> -> tensor<64x32x!tt.ptr<f32>, #blocked>
%68 = tt.addptr %67, %66 : tensor<64x32x!tt.ptr<f32>, #blocked>, tensor<64x32xi64, #blocked>
%69 = ttg.convert_layout %63 : tensor<64x32xf32, #mma> -> tensor<64x32xf32, #blocked>
tt.store %68, %69 : tensor<64x32x!tt.ptr<f32>, #blocked>
tt.return
}
} // end module
bench/specimen/artifacts/triton-ttgir.mlir
; ==== module part 00 of 04 ====
; ModuleID = '__compute_module_part_00'
source_filename = "__compute_module"
target datalayout = "e-m:o-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-darwin25.5.0"
%XLA_CPU_KernelCallFrame = type { ptr, ptr, i64, ptr }
%XLA_CPU_KernelThreadDim = type { i64, i64, i64 }
%XLA_CPU_KernelThread = type { i64, i64, i64 }
%XLA_CPU_KernelArg = type { ptr, i64 }
@__llvmsplit_unnamed.1 = private unnamed_addr constant [4 x i8] c"\00\00\80\FF"
@__llvmsplit_unnamed.2 = private unnamed_addr constant [4 x i8] c"\F3\045>"
; Function Attrs: uwtable
define ptr @multiply_reduce_fusion(ptr %0) #0 {
%reduce_function_parameter_addresses = alloca ptr, i32 2, align 8
%reduce_function_return_value_addr = alloca float, align 4
%arg_addr4 = alloca float, align 4
%arg_addr = alloca float, align 4
%reduce.0.inner.invar_address.reduction_dim.1 = alloca i64, align 8
%accumulator_0 = alloca float, align 4
%multiply_reduce_fusion.invar_address.dim.0 = alloca i64, align 8
%tdims_gep = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 0
%tdims = load ptr, ptr %tdims_gep, align 8
%tdim_x_gep = getelementptr inbounds nuw %XLA_CPU_KernelThreadDim, ptr %tdims, i32 0, i32 0
%tdim_y_gep = getelementptr inbounds nuw %XLA_CPU_KernelThreadDim, ptr %tdims, i32 0, i32 1
%tdim_z_gep = getelementptr inbounds nuw %XLA_CPU_KernelThreadDim, ptr %tdims, i32 0, i32 2
%tdim_x = load i64, ptr %tdim_x_gep, align 4
%tdim_y = load i64, ptr %tdim_y_gep, align 4
%tdim_z = load i64, ptr %tdim_z_gep, align 4
%tid_gep = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 1
%tids = load ptr, ptr %tid_gep, align 8
%tid_x_gep = getelementptr inbounds nuw %XLA_CPU_KernelThread, ptr %tids, i32 0, i32 0
%tid_y_gep = getelementptr inbounds nuw %XLA_CPU_KernelThread, ptr %tids, i32 0, i32 1
%tid_z_gep = getelementptr inbounds nuw %XLA_CPU_KernelThread, ptr %tids, i32 0, i32 2
%tid_x = load i64, ptr %tid_x_gep, align 4
%tid_y = load i64, ptr %tid_y_gep, align 4
%tid_z = load i64, ptr %tid_z_gep, align 4
%args_gep = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 3
%args = load ptr, ptr %args_gep, align 8
%arg0_gep = getelementptr %XLA_CPU_KernelArg, ptr %args, i32 0, i32 0
%arg0 = load ptr, ptr %arg0_gep, align 8, !invariant.load !0, !dereferenceable !1, !align !2
%args_gep1 = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 3
%args2 = load ptr, ptr %args_gep1, align 8
%arg1_gep = getelementptr %XLA_CPU_KernelArg, ptr %args2, i32 1, i32 0
%arg1 = load ptr, ptr %arg1_gep, align 8, !invariant.load !0, !dereferenceable !2, !align !2
store i64 0, ptr %multiply_reduce_fusion.invar_address.dim.0, align 4
br label %multiply_reduce_fusion.loop_header.dim.0
multiply_reduce_fusion.loop_header.dim.0: ; preds = %reduce.0.inner.loop_exit.reduction_dim.1, %1
%multiply_reduce_fusion.indvar.dim.0 = load i64, ptr %multiply_reduce_fusion.invar_address.dim.0, align 4
%2 = icmp uge i64 %multiply_reduce_fusion.indvar.dim.0, 16
br i1 %2, label %multiply_reduce_fusion.loop_exit.dim.0, label %multiply_reduce_fusion.loop_body.dim.0
multiply_reduce_fusion.loop_body.dim.0: ; preds = %multiply_reduce_fusion.loop_header.dim.0
%constant.3 = load float, ptr @__llvmsplit_unnamed.1, align 4
store float %constant.3, ptr %accumulator_0, align 4
store i64 0, ptr %reduce.0.inner.invar_address.reduction_dim.1, align 4
br label %reduce.0.inner.loop_header.reduction_dim.1
reduce.0.inner.loop_header.reduction_dim.1: ; preds = %reduce.0.inner.loop_body.reduction_dim.1, %multiply_reduce_fusion.loop_body.dim.0
%reduce.0.inner.indvar.reduction_dim.1 = load i64, ptr %reduce.0.inner.invar_address.reduction_dim.1, align 4
%3 = icmp uge i64 %reduce.0.inner.indvar.reduction_dim.1, 16
br i1 %3, label %reduce.0.inner.loop_exit.reduction_dim.1, label %reduce.0.inner.loop_body.reduction_dim.1
reduce.0.inner.loop_body.reduction_dim.1: ; preds = %reduce.0.inner.loop_header.reduction_dim.1
%4 = load float, ptr %accumulator_0, align 4
%5 = getelementptr inbounds [16 x [16 x float]], ptr %arg0, i64 0, i64 %multiply_reduce_fusion.indvar.dim.0, i64 %reduce.0.inner.indvar.reduction_dim.1
%6 = load float, ptr %5, align 4, !invariant.load !0, !noalias !3
%constant.4 = load float, ptr @__llvmsplit_unnamed.2, align 4
%multiply.2 = fmul float %6, %constant.4
store float %4, ptr %arg_addr, align 4
store float %multiply.2, ptr %arg_addr4, align 4
%7 = getelementptr inbounds ptr, ptr %reduce_function_parameter_addresses, i64 0
store ptr %arg_addr, ptr %7, align 8
%8 = getelementptr inbounds ptr, ptr %reduce_function_parameter_addresses, i64 1
store ptr %arg_addr4, ptr %8, align 8
call void @reduce_function(ptr %reduce_function_return_value_addr, ptr null, ptr %reduce_function_parameter_addresses, ptr null, ptr null, ptr null)
%9 = load float, ptr %reduce_function_return_value_addr, align 4
store float %9, ptr %accumulator_0, align 4
%invar.inc3 = add nuw nsw i64 %reduce.0.inner.indvar.reduction_dim.1, 1
store i64 %invar.inc3, ptr %reduce.0.inner.invar_address.reduction_dim.1, align 4
br label %reduce.0.inner.loop_header.reduction_dim.1
reduce.0.inner.loop_exit.reduction_dim.1: ; preds = %reduce.0.inner.loop_header.reduction_dim.1
%10 = load float, ptr %accumulator_0, align 4
%11 = getelementptr inbounds [16 x float], ptr %arg1, i64 0, i64 %multiply_reduce_fusion.indvar.dim.0
store float %10, ptr %11, align 4, !alias.scope !3
%invar.inc = add nuw nsw i64 %multiply_reduce_fusion.indvar.dim.0, 1
store i64 %invar.inc, ptr %multiply_reduce_fusion.invar_address.dim.0, align 4
br label %multiply_reduce_fusion.loop_header.dim.0
multiply_reduce_fusion.loop_exit.dim.0: ; preds = %multiply_reduce_fusion.loop_header.dim.0
br label %return
return: ; preds = %multiply_reduce_fusion.loop_exit.dim.0
ret ptr null
}
; Function Attrs: alwaysinline uwtable
define internal void @reduce_function(ptr %retval, ptr noalias %run_options, ptr noalias %params, ptr noalias %buffer_table, ptr noalias %status, ptr noalias %prof_counters) #1 {
entry:
%maximum.20 = alloca float, align 4
%0 = getelementptr inbounds ptr, ptr %params, i64 0
%Arg_0.18 = load ptr, ptr %0, align 8, !dereferenceable !6, !align !6
%1 = getelementptr inbounds ptr, ptr %params, i64 1
%Arg_1.19 = load ptr, ptr %1, align 8, !dereferenceable !6, !align !6
%2 = load float, ptr %Arg_0.18, align 4, !alias.scope !7, !noalias !10
%3 = load float, ptr %Arg_1.19, align 4, !alias.scope !12, !noalias !10
%4 = call reassoc float @llvm.maximum.f32(float %2, float %3)
store float %4, ptr %maximum.20, align 4, !alias.scope !10
%load_ret_value = load float, ptr %maximum.20, align 4
store float %load_ret_value, ptr %retval, align 4
br label %return
return: ; preds = %entry
ret void
}
; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
declare float @llvm.maximum.f32(float, float) #2
attributes #0 = { uwtable "frame-pointer"="all" "prefer-vector-width"="256" }
attributes #1 = { alwaysinline uwtable "denormal-fp-math"="preserve-sign" "frame-pointer"="none" }
attributes #2 = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
!0 = !{}
!1 = !{i64 1024}
!2 = !{i64 64}
!3 = !{!4}
!4 = !{!"result slice: {index:12, offset:0, size:64}", !5}
!5 = !{!"XLA host kernel multiply_reduce_fusion AA domain"}
!6 = !{i64 4}
!7 = !{!8}
!8 = !{!"buffer: {index:7, offset:0, size:4}", !9}
!9 = !{!"XLA global AA domain"}
!10 = !{!11}
!11 = !{!"buffer: {index:9, offset:0, size:4}", !9}
!12 = !{!13}
!13 = !{!"buffer: {index:8, offset:0, size:4}", !9}
; ==== module part 01 of 04 ====
; ModuleID = '__compute_module_part_01'
source_filename = "__compute_module"
target datalayout = "e-m:o-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-darwin25.5.0"
%XLA_CPU_KernelCallFrame = type { ptr, ptr, i64, ptr }
%XLA_CPU_KernelThreadDim = type { i64, i64, i64 }
%XLA_CPU_KernelThread = type { i64, i64, i64 }
%XLA_CPU_KernelArg = type { ptr, i64 }
@__llvmsplit_unnamed.3 = private unnamed_addr constant [4 x i8] c"\00\00\80\FF"
@__llvmsplit_unnamed.4 = private unnamed_addr constant [4 x i8] c"\F3\045>"
; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
declare float @llvm.maximum.f32(float, float) #0
; Function Attrs: uwtable
define ptr @subtract_exponential_fusion(ptr %0) #1 {
%subtract_exponential_fusion.invar_address.dim.1 = alloca i64, align 8
%subtract_exponential_fusion.invar_address.dim.0 = alloca i64, align 8
%tdims_gep = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 0
%tdims = load ptr, ptr %tdims_gep, align 8
%tdim_x_gep = getelementptr inbounds nuw %XLA_CPU_KernelThreadDim, ptr %tdims, i32 0, i32 0
%tdim_y_gep = getelementptr inbounds nuw %XLA_CPU_KernelThreadDim, ptr %tdims, i32 0, i32 1
%tdim_z_gep = getelementptr inbounds nuw %XLA_CPU_KernelThreadDim, ptr %tdims, i32 0, i32 2
%tdim_x = load i64, ptr %tdim_x_gep, align 4
%tdim_y = load i64, ptr %tdim_y_gep, align 4
%tdim_z = load i64, ptr %tdim_z_gep, align 4
%tid_gep = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 1
%tids = load ptr, ptr %tid_gep, align 8
%tid_x_gep = getelementptr inbounds nuw %XLA_CPU_KernelThread, ptr %tids, i32 0, i32 0
%tid_y_gep = getelementptr inbounds nuw %XLA_CPU_KernelThread, ptr %tids, i32 0, i32 1
%tid_z_gep = getelementptr inbounds nuw %XLA_CPU_KernelThread, ptr %tids, i32 0, i32 2
%tid_x = load i64, ptr %tid_x_gep, align 4
%tid_y = load i64, ptr %tid_y_gep, align 4
%tid_z = load i64, ptr %tid_z_gep, align 4
%args_gep = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 3
%args = load ptr, ptr %args_gep, align 8
%arg0_gep = getelementptr %XLA_CPU_KernelArg, ptr %args, i32 0, i32 0
%arg0 = load ptr, ptr %arg0_gep, align 8, !invariant.load !0, !dereferenceable !1, !align !1
%args_gep1 = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 3
%args2 = load ptr, ptr %args_gep1, align 8
%arg1_gep = getelementptr %XLA_CPU_KernelArg, ptr %args2, i32 1, i32 0
%arg1 = load ptr, ptr %arg1_gep, align 8, !invariant.load !0, !dereferenceable !2, !align !1
%args_gep3 = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 3
%args4 = load ptr, ptr %args_gep3, align 8
%arg2_gep = getelementptr %XLA_CPU_KernelArg, ptr %args4, i32 2, i32 0
%arg2 = load ptr, ptr %arg2_gep, align 8, !invariant.load !0, !dereferenceable !2, !align !1
store i64 0, ptr %subtract_exponential_fusion.invar_address.dim.0, align 4
br label %subtract_exponential_fusion.loop_header.dim.0
subtract_exponential_fusion.loop_header.dim.0: ; preds = %subtract_exponential_fusion.loop_exit.dim.1, %1
%subtract_exponential_fusion.indvar.dim.0 = load i64, ptr %subtract_exponential_fusion.invar_address.dim.0, align 4
%2 = icmp uge i64 %subtract_exponential_fusion.indvar.dim.0, 16
br i1 %2, label %subtract_exponential_fusion.loop_exit.dim.0, label %subtract_exponential_fusion.loop_body.dim.0
subtract_exponential_fusion.loop_body.dim.0: ; preds = %subtract_exponential_fusion.loop_header.dim.0
store i64 0, ptr %subtract_exponential_fusion.invar_address.dim.1, align 4
br label %subtract_exponential_fusion.loop_header.dim.1
subtract_exponential_fusion.loop_header.dim.1: ; preds = %subtract_exponential_fusion.loop_body.dim.1, %subtract_exponential_fusion.loop_body.dim.0
%subtract_exponential_fusion.indvar.dim.1 = load i64, ptr %subtract_exponential_fusion.invar_address.dim.1, align 4
%3 = icmp uge i64 %subtract_exponential_fusion.indvar.dim.1, 16
br i1 %3, label %subtract_exponential_fusion.loop_exit.dim.1, label %subtract_exponential_fusion.loop_body.dim.1
subtract_exponential_fusion.loop_body.dim.1: ; preds = %subtract_exponential_fusion.loop_header.dim.1
%4 = getelementptr inbounds [16 x [16 x float]], ptr %arg1, i64 0, i64 %subtract_exponential_fusion.indvar.dim.0, i64 %subtract_exponential_fusion.indvar.dim.1
%5 = load float, ptr %4, align 4, !invariant.load !0, !noalias !3
%constant.2 = load float, ptr @__llvmsplit_unnamed.4, align 4
%multiply.1 = fmul float %5, %constant.2
%6 = getelementptr inbounds [16 x float], ptr %arg0, i64 0, i64 %subtract_exponential_fusion.indvar.dim.0
%7 = load float, ptr %6, align 4, !invariant.load !0, !noalias !3
%constant.1 = load float, ptr @__llvmsplit_unnamed.3, align 4
%8 = call float @llvm.maximum.f32(float %7, float %constant.1)
%subtract.0 = fsub float %multiply.1, %8
%9 = call float @llvm.exp.f32(float %subtract.0)
%10 = getelementptr inbounds [16 x [16 x float]], ptr %arg2, i64 0, i64 %subtract_exponential_fusion.indvar.dim.0, i64 %subtract_exponential_fusion.indvar.dim.1
store float %9, ptr %10, align 4, !alias.scope !3
%invar.inc5 = add nuw nsw i64 %subtract_exponential_fusion.indvar.dim.1, 1
store i64 %invar.inc5, ptr %subtract_exponential_fusion.invar_address.dim.1, align 4
br label %subtract_exponential_fusion.loop_header.dim.1
subtract_exponential_fusion.loop_exit.dim.1: ; preds = %subtract_exponential_fusion.loop_header.dim.1
%invar.inc = add nuw nsw i64 %subtract_exponential_fusion.indvar.dim.0, 1
store i64 %invar.inc, ptr %subtract_exponential_fusion.invar_address.dim.0, align 4
br label %subtract_exponential_fusion.loop_header.dim.0
subtract_exponential_fusion.loop_exit.dim.0: ; preds = %subtract_exponential_fusion.loop_header.dim.0
br label %return
return: ; preds = %subtract_exponential_fusion.loop_exit.dim.0
ret ptr null
}
; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
declare float @llvm.exp.f32(float) #0
attributes #0 = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
attributes #1 = { uwtable "frame-pointer"="all" "prefer-vector-width"="256" }
!0 = !{}
!1 = !{i64 64}
!2 = !{i64 1024}
!3 = !{!4}
!4 = !{!"result slice: {index:3, offset:0, size:1024}", !5}
!5 = !{!"XLA host kernel subtract_exponential_fusion AA domain"}
; ==== module part 02 of 04 ====
; ModuleID = '__compute_module_part_02'
source_filename = "__compute_module"
target datalayout = "e-m:o-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-darwin25.5.0"
%XLA_CPU_KernelCallFrame = type { ptr, ptr, i64, ptr }
%XLA_CPU_KernelThreadDim = type { i64, i64, i64 }
%XLA_CPU_KernelThread = type { i64, i64, i64 }
%XLA_CPU_KernelArg = type { ptr, i64 }
; Function Attrs: uwtable
define ptr @reduce.33(ptr %0) #0 {
%reduce_function_parameter_addresses = alloca ptr, i32 2, align 8
%reduce_function_return_value_addr = alloca float, align 4
%arg_addr6 = alloca float, align 4
%arg_addr = alloca float, align 4
%reduce.33.inner.invar_address.reduction_dim.1 = alloca i64, align 8
%accumulator_0 = alloca float, align 4
%reduce.33.invar_address.dim.0 = alloca i64, align 8
%tdims_gep = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 0
%tdims = load ptr, ptr %tdims_gep, align 8
%tdim_x_gep = getelementptr inbounds nuw %XLA_CPU_KernelThreadDim, ptr %tdims, i32 0, i32 0
%tdim_y_gep = getelementptr inbounds nuw %XLA_CPU_KernelThreadDim, ptr %tdims, i32 0, i32 1
%tdim_z_gep = getelementptr inbounds nuw %XLA_CPU_KernelThreadDim, ptr %tdims, i32 0, i32 2
%tdim_x = load i64, ptr %tdim_x_gep, align 4
%tdim_y = load i64, ptr %tdim_y_gep, align 4
%tdim_z = load i64, ptr %tdim_z_gep, align 4
%tid_gep = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 1
%tids = load ptr, ptr %tid_gep, align 8
%tid_x_gep = getelementptr inbounds nuw %XLA_CPU_KernelThread, ptr %tids, i32 0, i32 0
%tid_y_gep = getelementptr inbounds nuw %XLA_CPU_KernelThread, ptr %tids, i32 0, i32 1
%tid_z_gep = getelementptr inbounds nuw %XLA_CPU_KernelThread, ptr %tids, i32 0, i32 2
%tid_x = load i64, ptr %tid_x_gep, align 4
%tid_y = load i64, ptr %tid_y_gep, align 4
%tid_z = load i64, ptr %tid_z_gep, align 4
%args_gep = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 3
%args = load ptr, ptr %args_gep, align 8
%arg0_gep = getelementptr %XLA_CPU_KernelArg, ptr %args, i32 0, i32 0
%arg0 = load ptr, ptr %arg0_gep, align 8, !invariant.load !0, !dereferenceable !1, !align !2
%args_gep1 = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 3
%args2 = load ptr, ptr %args_gep1, align 8
%arg1_gep = getelementptr %XLA_CPU_KernelArg, ptr %args2, i32 1, i32 0
%arg1 = load ptr, ptr %arg1_gep, align 8, !invariant.load !0, !dereferenceable !3, !align !2
%args_gep3 = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 3
%args4 = load ptr, ptr %args_gep3, align 8
%arg2_gep = getelementptr %XLA_CPU_KernelArg, ptr %args4, i32 2, i32 0
%arg2 = load ptr, ptr %arg2_gep, align 8, !invariant.load !0, !dereferenceable !2, !align !2
store i64 0, ptr %reduce.33.invar_address.dim.0, align 4
br label %reduce.33.loop_header.dim.0
reduce.33.loop_header.dim.0: ; preds = %reduce.33.inner.loop_exit.reduction_dim.1, %1
%reduce.33.indvar.dim.0 = load i64, ptr %reduce.33.invar_address.dim.0, align 4
%2 = icmp uge i64 %reduce.33.indvar.dim.0, 16
br i1 %2, label %reduce.33.loop_exit.dim.0, label %reduce.33.loop_body.dim.0
reduce.33.loop_body.dim.0: ; preds = %reduce.33.loop_header.dim.0
%3 = load float, ptr %arg1, align 4, !invariant.load !0, !noalias !4
store float %3, ptr %accumulator_0, align 4
store i64 0, ptr %reduce.33.inner.invar_address.reduction_dim.1, align 4
br label %reduce.33.inner.loop_header.reduction_dim.1
reduce.33.inner.loop_header.reduction_dim.1: ; preds = %reduce.33.inner.loop_body.reduction_dim.1, %reduce.33.loop_body.dim.0
%reduce.33.inner.indvar.reduction_dim.1 = load i64, ptr %reduce.33.inner.invar_address.reduction_dim.1, align 4
%4 = icmp uge i64 %reduce.33.inner.indvar.reduction_dim.1, 16
br i1 %4, label %reduce.33.inner.loop_exit.reduction_dim.1, label %reduce.33.inner.loop_body.reduction_dim.1
reduce.33.inner.loop_body.reduction_dim.1: ; preds = %reduce.33.inner.loop_header.reduction_dim.1
%5 = load float, ptr %accumulator_0, align 4
%6 = getelementptr inbounds [16 x [16 x float]], ptr %arg0, i64 0, i64 %reduce.33.indvar.dim.0, i64 %reduce.33.inner.indvar.reduction_dim.1
%7 = load float, ptr %6, align 4, !invariant.load !0, !noalias !4
store float %5, ptr %arg_addr, align 4
store float %7, ptr %arg_addr6, align 4
%8 = getelementptr inbounds ptr, ptr %reduce_function_parameter_addresses, i64 0
store ptr %arg_addr, ptr %8, align 8
%9 = getelementptr inbounds ptr, ptr %reduce_function_parameter_addresses, i64 1
store ptr %arg_addr6, ptr %9, align 8
call void @reduce_function__1(ptr %reduce_function_return_value_addr, ptr null, ptr %reduce_function_parameter_addresses, ptr null, ptr null, ptr null)
%10 = load float, ptr %reduce_function_return_value_addr, align 4
store float %10, ptr %accumulator_0, align 4
%invar.inc5 = add nuw nsw i64 %reduce.33.inner.indvar.reduction_dim.1, 1
store i64 %invar.inc5, ptr %reduce.33.inner.invar_address.reduction_dim.1, align 4
br label %reduce.33.inner.loop_header.reduction_dim.1
reduce.33.inner.loop_exit.reduction_dim.1: ; preds = %reduce.33.inner.loop_header.reduction_dim.1
%11 = load float, ptr %accumulator_0, align 4
%12 = getelementptr inbounds [16 x float], ptr %arg2, i64 0, i64 %reduce.33.indvar.dim.0
store float %11, ptr %12, align 4, !alias.scope !4
%invar.inc = add nuw nsw i64 %reduce.33.indvar.dim.0, 1
store i64 %invar.inc, ptr %reduce.33.invar_address.dim.0, align 4
br label %reduce.33.loop_header.dim.0
reduce.33.loop_exit.dim.0: ; preds = %reduce.33.loop_header.dim.0
br label %return
return: ; preds = %reduce.33.loop_exit.dim.0
ret ptr null
}
; Function Attrs: alwaysinline uwtable
define internal void @reduce_function__1(ptr %retval, ptr noalias %run_options, ptr noalias %params, ptr noalias %buffer_table, ptr noalias %status, ptr noalias %prof_counters) #1 {
entry:
%add.32 = alloca float, align 4
%0 = getelementptr inbounds ptr, ptr %params, i64 0
%Arg_0.30 = load ptr, ptr %0, align 8, !dereferenceable !3, !align !3
%1 = getelementptr inbounds ptr, ptr %params, i64 1
%Arg_1.31 = load ptr, ptr %1, align 8, !dereferenceable !3, !align !3
%2 = load float, ptr %Arg_0.30, align 4, !alias.scope !7, !noalias !10
%3 = load float, ptr %Arg_1.31, align 4, !alias.scope !12, !noalias !10
%add.321 = fadd reassoc float %2, %3
store float %add.321, ptr %add.32, align 4, !alias.scope !10
%load_ret_value = load float, ptr %add.32, align 4
store float %load_ret_value, ptr %retval, align 4
br label %return
return: ; preds = %entry
ret void
}
attributes #0 = { uwtable "frame-pointer"="all" "prefer-vector-width"="256" }
attributes #1 = { alwaysinline uwtable "denormal-fp-math"="preserve-sign" "frame-pointer"="none" }
!0 = !{}
!1 = !{i64 1024}
!2 = !{i64 64}
!3 = !{i64 4}
!4 = !{!5}
!5 = !{!"result slice: {index:12, offset:0, size:64}", !6}
!6 = !{!"XLA host kernel reduce.33 AA domain"}
!7 = !{!8}
!8 = !{!"buffer: {index:10, offset:0, size:4}", !9}
!9 = !{!"XLA global AA domain"}
!10 = !{!11}
!11 = !{!"buffer: {index:6, offset:0, size:4}", !9}
!12 = !{!13}
!13 = !{!"buffer: {index:11, offset:0, size:4}", !9}
; ==== module part 03 of 04 ====
; ModuleID = '__compute_module_part_03'
source_filename = "__compute_module"
target datalayout = "e-m:o-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-darwin25.5.0"
%XLA_CPU_KernelCallFrame = type { ptr, ptr, i64, ptr }
%XLA_CPU_KernelThreadDim = type { i64, i64, i64 }
%XLA_CPU_KernelThread = type { i64, i64, i64 }
%XLA_CPU_KernelArg = type { ptr, i64 }
; Function Attrs: uwtable
define ptr @broadcast_divide_fusion(ptr %0) #0 {
%broadcast_divide_fusion.invar_address.dim.1 = alloca i64, align 8
%broadcast_divide_fusion.invar_address.dim.0 = alloca i64, align 8
%tdims_gep = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 0
%tdims = load ptr, ptr %tdims_gep, align 8
%tdim_x_gep = getelementptr inbounds nuw %XLA_CPU_KernelThreadDim, ptr %tdims, i32 0, i32 0
%tdim_y_gep = getelementptr inbounds nuw %XLA_CPU_KernelThreadDim, ptr %tdims, i32 0, i32 1
%tdim_z_gep = getelementptr inbounds nuw %XLA_CPU_KernelThreadDim, ptr %tdims, i32 0, i32 2
%tdim_x = load i64, ptr %tdim_x_gep, align 4
%tdim_y = load i64, ptr %tdim_y_gep, align 4
%tdim_z = load i64, ptr %tdim_z_gep, align 4
%tid_gep = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 1
%tids = load ptr, ptr %tid_gep, align 8
%tid_x_gep = getelementptr inbounds nuw %XLA_CPU_KernelThread, ptr %tids, i32 0, i32 0
%tid_y_gep = getelementptr inbounds nuw %XLA_CPU_KernelThread, ptr %tids, i32 0, i32 1
%tid_z_gep = getelementptr inbounds nuw %XLA_CPU_KernelThread, ptr %tids, i32 0, i32 2
%tid_x = load i64, ptr %tid_x_gep, align 4
%tid_y = load i64, ptr %tid_y_gep, align 4
%tid_z = load i64, ptr %tid_z_gep, align 4
%args_gep = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 3
%args = load ptr, ptr %args_gep, align 8
%arg0_gep = getelementptr %XLA_CPU_KernelArg, ptr %args, i32 0, i32 0
%arg0 = load ptr, ptr %arg0_gep, align 8, !invariant.load !0, !dereferenceable !1, !align !2
%args_gep1 = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 3
%args2 = load ptr, ptr %args_gep1, align 8
%arg1_gep = getelementptr %XLA_CPU_KernelArg, ptr %args2, i32 1, i32 0
%arg1 = load ptr, ptr %arg1_gep, align 8, !invariant.load !0, !dereferenceable !2, !align !2
%args_gep3 = getelementptr inbounds nuw %XLA_CPU_KernelCallFrame, ptr %0, i32 0, i32 3
%args4 = load ptr, ptr %args_gep3, align 8
%arg2_gep = getelementptr %XLA_CPU_KernelArg, ptr %args4, i32 2, i32 0
%arg2 = load ptr, ptr %arg2_gep, align 8, !invariant.load !0, !dereferenceable !1, !align !2
store i64 0, ptr %broadcast_divide_fusion.invar_address.dim.0, align 4
br label %broadcast_divide_fusion.loop_header.dim.0
broadcast_divide_fusion.loop_header.dim.0: ; preds = %broadcast_divide_fusion.loop_exit.dim.1, %1
%broadcast_divide_fusion.indvar.dim.0 = load i64, ptr %broadcast_divide_fusion.invar_address.dim.0, align 4
%2 = icmp uge i64 %broadcast_divide_fusion.indvar.dim.0, 16
br i1 %2, label %broadcast_divide_fusion.loop_exit.dim.0, label %broadcast_divide_fusion.loop_body.dim.0
broadcast_divide_fusion.loop_body.dim.0: ; preds = %broadcast_divide_fusion.loop_header.dim.0
store i64 0, ptr %broadcast_divide_fusion.invar_address.dim.1, align 4
br label %broadcast_divide_fusion.loop_header.dim.1
broadcast_divide_fusion.loop_header.dim.1: ; preds = %broadcast_divide_fusion.loop_body.dim.1, %broadcast_divide_fusion.loop_body.dim.0
%broadcast_divide_fusion.indvar.dim.1 = load i64, ptr %broadcast_divide_fusion.invar_address.dim.1, align 4
%3 = icmp uge i64 %broadcast_divide_fusion.indvar.dim.1, 16
br i1 %3, label %broadcast_divide_fusion.loop_exit.dim.1, label %broadcast_divide_fusion.loop_body.dim.1
broadcast_divide_fusion.loop_body.dim.1: ; preds = %broadcast_divide_fusion.loop_header.dim.1
%4 = getelementptr inbounds [16 x [16 x float]], ptr %arg0, i64 0, i64 %broadcast_divide_fusion.indvar.dim.0, i64 %broadcast_divide_fusion.indvar.dim.1
%5 = load float, ptr %4, align 4, !invariant.load !0, !noalias !3
%6 = getelementptr inbounds [16 x float], ptr %arg1, i64 0, i64 %broadcast_divide_fusion.indvar.dim.0
%7 = load float, ptr %6, align 4, !invariant.load !0, !noalias !3
%divide.0 = fdiv float %5, %7
%8 = getelementptr inbounds [16 x [16 x float]], ptr %arg2, i64 0, i64 %broadcast_divide_fusion.indvar.dim.0, i64 %broadcast_divide_fusion.indvar.dim.1
store float %divide.0, ptr %8, align 4, !alias.scope !3
%invar.inc5 = add nuw nsw i64 %broadcast_divide_fusion.indvar.dim.1, 1
store i64 %invar.inc5, ptr %broadcast_divide_fusion.invar_address.dim.1, align 4
br label %broadcast_divide_fusion.loop_header.dim.1
broadcast_divide_fusion.loop_exit.dim.1: ; preds = %broadcast_divide_fusion.loop_header.dim.1
%invar.inc = add nuw nsw i64 %broadcast_divide_fusion.indvar.dim.0, 1
store i64 %invar.inc, ptr %broadcast_divide_fusion.invar_address.dim.0, align 4
br label %broadcast_divide_fusion.loop_header.dim.0
broadcast_divide_fusion.loop_exit.dim.0: ; preds = %broadcast_divide_fusion.loop_header.dim.0
br label %return
return: ; preds = %broadcast_divide_fusion.loop_exit.dim.0
ret ptr null
}
attributes #0 = { uwtable "frame-pointer"="all" "prefer-vector-width"="256" }
!0 = !{}
!1 = !{i64 1024}
!2 = !{i64 64}
!3 = !{!4}
!4 = !{!"result slice: {index:12, offset:2048, size:1024}", !5}
!5 = !{!"XLA host kernel broadcast_divide_fusion AA domain"}bench/specimen/artifacts/llvm-ir-no-opt.ll
; ==== module part 00 of 04 ====
; ModuleID = '__compute_module_part_00'
source_filename = "__compute_module"
target datalayout = "e-m:o-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-darwin25.5.0"
; Function Attrs: nofree norecurse nosync nounwind memory(readwrite, inaccessiblemem: none) uwtable
define noalias noundef ptr @multiply_reduce_fusion(ptr nocapture readonly %0) local_unnamed_addr #0 {
%args_gep = getelementptr inbounds nuw i8, ptr %0, i64 24
%args = load ptr, ptr %args_gep, align 8
%arg0 = load ptr, ptr %args, align 8, !invariant.load !0, !dereferenceable !1, !align !2
%arg1_gep = getelementptr i8, ptr %args, i64 16
%arg1 = load ptr, ptr %arg1_gep, align 8, !invariant.load !0, !dereferenceable !2, !align !2
br label %middle.block
middle.block: ; preds = %middle.block, %1
%multiply_reduce_fusion.invar_address.dim.0.05 = phi i64 [ 0, %1 ], [ %invar.inc, %middle.block ]
%2 = getelementptr inbounds nuw [16 x [16 x float]], ptr %arg0, i64 0, i64 %multiply_reduce_fusion.invar_address.dim.0.05, i64 0
%3 = getelementptr inbounds nuw i8, ptr %2, i64 32
%wide.load8 = load <8 x float>, ptr %3, align 32, !invariant.load !0, !noalias !3
%4 = fmul <8 x float> %wide.load8, splat (float 0x3FC6A09E60000000)
%wide.load = load <8 x float>, ptr %2, align 64, !invariant.load !0, !noalias !3
%5 = fmul <8 x float> %wide.load, splat (float 0x3FC6A09E60000000)
%rdx.minmax = tail call reassoc <8 x float> @llvm.maximum.v8f32(<8 x float> %5, <8 x float> %4)
%6 = tail call reassoc float @llvm.vector.reduce.fmaximum.v8f32(<8 x float> %rdx.minmax)
%7 = getelementptr inbounds nuw [16 x float], ptr %arg1, i64 0, i64 %multiply_reduce_fusion.invar_address.dim.0.05
store float %6, ptr %7, align 4, !alias.scope !3
%invar.inc = add nuw nsw i64 %multiply_reduce_fusion.invar_address.dim.0.05, 1
%exitcond6 = icmp eq i64 %invar.inc, 16
br i1 %exitcond6, label %return, label %middle.block
return: ; preds = %middle.block
ret ptr null
}
; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
declare <8 x float> @llvm.maximum.v8f32(<8 x float>, <8 x float>) #1
; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
declare float @llvm.vector.reduce.fmaximum.v8f32(<8 x float>) #1
attributes #0 = { nofree norecurse nosync nounwind memory(readwrite, inaccessiblemem: none) uwtable "frame-pointer"="all" "prefer-vector-width"="256" }
attributes #1 = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
!0 = !{}
!1 = !{i64 1024}
!2 = !{i64 64}
!3 = !{!4}
!4 = !{!"result slice: {index:12, offset:0, size:64}", !5}
!5 = !{!"XLA host kernel multiply_reduce_fusion AA domain"}
; ==== module part 01 of 04 ====
; ModuleID = '__compute_module_part_01'
source_filename = "__compute_module"
target datalayout = "e-m:o-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-darwin25.5.0"
; Function Attrs: nofree norecurse nosync nounwind memory(readwrite, inaccessiblemem: none) uwtable
define noalias noundef ptr @subtract_exponential_fusion(ptr nocapture readonly %0) local_unnamed_addr #0 {
%args_gep = getelementptr inbounds nuw i8, ptr %0, i64 24
%args = load ptr, ptr %args_gep, align 8
%arg0 = load ptr, ptr %args, align 8, !invariant.load !0, !dereferenceable !1, !align !1
%arg1_gep = getelementptr i8, ptr %args, i64 16
%arg1 = load ptr, ptr %arg1_gep, align 8, !invariant.load !0, !dereferenceable !2, !align !1
%arg2_gep = getelementptr i8, ptr %args, i64 32
%arg2 = load ptr, ptr %arg2_gep, align 8, !invariant.load !0, !dereferenceable !2, !align !1
br label %vector.ph
vector.ph: ; preds = %vector.ph, %1
%subtract_exponential_fusion.invar_address.dim.0.02 = phi i64 [ 0, %1 ], [ %invar.inc, %vector.ph ]
%2 = getelementptr inbounds nuw [16 x float], ptr %arg0, i64 0, i64 %subtract_exponential_fusion.invar_address.dim.0.02
%3 = load float, ptr %2, align 4, !invariant.load !0, !noalias !3
%broadcast.splatinsert = insertelement <8 x float> poison, float %3, i64 0
%broadcast.splat = shufflevector <8 x float> %broadcast.splatinsert, <8 x float> poison, <8 x i32> zeroinitializer
%4 = getelementptr inbounds nuw [16 x [16 x float]], ptr %arg1, i64 0, i64 %subtract_exponential_fusion.invar_address.dim.0.02, i64 0
%5 = getelementptr inbounds nuw i8, ptr %4, i64 32
%wide.load = load <8 x float>, ptr %4, align 64, !invariant.load !0, !noalias !3
%wide.load4 = load <8 x float>, ptr %5, align 32, !invariant.load !0, !noalias !3
%6 = fmul <8 x float> %wide.load, splat (float 0x3FC6A09E60000000)
%7 = fmul <8 x float> %wide.load4, splat (float 0x3FC6A09E60000000)
%8 = fsub <8 x float> %6, %broadcast.splat
%9 = fsub <8 x float> %7, %broadcast.splat
%10 = fcmp uge <8 x float> %8, splat (float 0xC055F33340000000)
%11 = select <8 x i1> %10, <8 x float> %8, <8 x float> splat (float 0xC055F33340000000)
%12 = fcmp ule <8 x float> %11, splat (float 0x4056333340000000)
%13 = select <8 x i1> %12, <8 x float> %11, <8 x float> splat (float 0x4056333340000000)
%exp_f32.i5 = fmul <8 x float> %13, splat (float 0x3FF7154760000000)
%exp_f321.i6 = fadd <8 x float> splat (float 5.000000e-01), %exp_f32.i5
%14 = call <8 x float> @llvm.floor.v8f32(<8 x float> %exp_f321.i6)
%15 = fcmp uge <8 x float> %14, splat (float -1.270000e+02)
%16 = select <8 x i1> %15, <8 x float> %14, <8 x float> splat (float -1.270000e+02)
%17 = fcmp ule <8 x float> %16, splat (float 1.270000e+02)
%18 = select <8 x i1> %17, <8 x float> %16, <8 x float> splat (float 1.270000e+02)
%exp_f322.i7 = fmul <8 x float> splat (float 0x3FE6300000000000), %18
%19 = fsub <8 x float> %13, %exp_f322.i7
%exp_f323.i8 = fmul <8 x float> splat (float 0xBF2BD01060000000), %18
%20 = fsub <8 x float> %19, %exp_f323.i8
%exp_f324.i9 = fmul <8 x float> %20, splat (float 0x3F2A0D2CE0000000)
%exp_f325.i10 = fadd <8 x float> splat (float 0x3F56E879C0000000), %exp_f324.i9
%exp_f326.i11 = fmul <8 x float> %exp_f325.i10, %20
%exp_f327.i12 = fadd <8 x float> splat (float 0x3F81112100000000), %exp_f326.i11
%exp_f328.i13 = fmul <8 x float> %exp_f327.i12, %20
%exp_f329.i14 = fadd <8 x float> splat (float 0x3FA5553820000000), %exp_f328.i13
%exp_f3210.i15 = fmul <8 x float> %exp_f329.i14, %20
%exp_f3211.i16 = fadd <8 x float> splat (float 0x3FC5555540000000), %exp_f3210.i15
%exp_f3212.i17 = fmul <8 x float> %exp_f3211.i16, %20
%exp_f3213.i18 = fadd <8 x float> splat (float 5.000000e-01), %exp_f3212.i17
%exp_f3214.i19 = fmul <8 x float> %20, %20
%exp_f3215.i20 = fmul <8 x float> %exp_f3213.i18, %exp_f3214.i19
%exp_f3216.i21 = fadd <8 x float> %20, %exp_f3215.i20
%exp_f3217.i22 = fadd <8 x float> splat (float 1.000000e+00), %exp_f3216.i21
%21 = fptosi <8 x float> %18 to <8 x i32>
%22 = add <8 x i32> %21, splat (i32 127)
%23 = shl <8 x i32> %22, splat (i32 23)
%24 = bitcast <8 x i32> %23 to <8 x float>
%exp_f3218.i23 = fmul <8 x float> %exp_f3217.i22, %24
%25 = fcmp uge <8 x float> %9, splat (float 0xC055F33340000000)
%26 = select <8 x i1> %25, <8 x float> %9, <8 x float> splat (float 0xC055F33340000000)
%27 = fcmp ule <8 x float> %26, splat (float 0x4056333340000000)
%28 = select <8 x i1> %27, <8 x float> %26, <8 x float> splat (float 0x4056333340000000)
%exp_f32.i = fmul <8 x float> %28, splat (float 0x3FF7154760000000)
%exp_f321.i = fadd <8 x float> splat (float 5.000000e-01), %exp_f32.i
%29 = call <8 x float> @llvm.floor.v8f32(<8 x float> %exp_f321.i)
%30 = fcmp uge <8 x float> %29, splat (float -1.270000e+02)
%31 = select <8 x i1> %30, <8 x float> %29, <8 x float> splat (float -1.270000e+02)
%32 = fcmp ule <8 x float> %31, splat (float 1.270000e+02)
%33 = select <8 x i1> %32, <8 x float> %31, <8 x float> splat (float 1.270000e+02)
%exp_f322.i = fmul <8 x float> splat (float 0x3FE6300000000000), %33
%34 = fsub <8 x float> %28, %exp_f322.i
%exp_f323.i = fmul <8 x float> splat (float 0xBF2BD01060000000), %33
%35 = fsub <8 x float> %34, %exp_f323.i
%exp_f324.i = fmul <8 x float> %35, splat (float 0x3F2A0D2CE0000000)
%exp_f325.i = fadd <8 x float> splat (float 0x3F56E879C0000000), %exp_f324.i
%exp_f326.i = fmul <8 x float> %exp_f325.i, %35
%exp_f327.i = fadd <8 x float> splat (float 0x3F81112100000000), %exp_f326.i
%exp_f328.i = fmul <8 x float> %exp_f327.i, %35
%exp_f329.i = fadd <8 x float> splat (float 0x3FA5553820000000), %exp_f328.i
%exp_f3210.i = fmul <8 x float> %exp_f329.i, %35
%exp_f3211.i = fadd <8 x float> splat (float 0x3FC5555540000000), %exp_f3210.i
%exp_f3212.i = fmul <8 x float> %exp_f3211.i, %35
%exp_f3213.i = fadd <8 x float> splat (float 5.000000e-01), %exp_f3212.i
%exp_f3214.i = fmul <8 x float> %35, %35
%exp_f3215.i = fmul <8 x float> %exp_f3213.i, %exp_f3214.i
%exp_f3216.i = fadd <8 x float> %35, %exp_f3215.i
%exp_f3217.i = fadd <8 x float> splat (float 1.000000e+00), %exp_f3216.i
%36 = fptosi <8 x float> %33 to <8 x i32>
%37 = add <8 x i32> %36, splat (i32 127)
%38 = shl <8 x i32> %37, splat (i32 23)
%39 = bitcast <8 x i32> %38 to <8 x float>
%exp_f3218.i = fmul <8 x float> %exp_f3217.i, %39
%40 = getelementptr inbounds nuw [16 x [16 x float]], ptr %arg2, i64 0, i64 %subtract_exponential_fusion.invar_address.dim.0.02, i64 0
%41 = getelementptr inbounds nuw i8, ptr %40, i64 32
store <8 x float> %exp_f3218.i23, ptr %40, align 64, !alias.scope !3
store <8 x float> %exp_f3218.i, ptr %41, align 32, !alias.scope !3
%invar.inc = add nuw nsw i64 %subtract_exponential_fusion.invar_address.dim.0.02, 1
%exitcond3 = icmp eq i64 %invar.inc, 16
br i1 %exitcond3, label %return, label %vector.ph
return: ; preds = %vector.ph
ret ptr null
}
; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
declare <4 x float> @llvm.floor.v4f32(<4 x float>) #1
; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
declare <8 x float> @llvm.floor.v8f32(<8 x float>) #1
; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
declare <16 x float> @llvm.floor.v16f32(<16 x float>) #1
attributes #0 = { nofree norecurse nosync nounwind memory(readwrite, inaccessiblemem: none) uwtable "frame-pointer"="all" "prefer-vector-width"="256" }
attributes #1 = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
!0 = !{}
!1 = !{i64 64}
!2 = !{i64 1024}
!3 = !{!4}
!4 = !{!"result slice: {index:3, offset:0, size:1024}", !5}
!5 = !{!"XLA host kernel subtract_exponential_fusion AA domain"}
; ==== module part 02 of 04 ====
; ModuleID = '__compute_module_part_02'
source_filename = "__compute_module"
target datalayout = "e-m:o-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-darwin25.5.0"
; Function Attrs: nofree norecurse nosync nounwind memory(readwrite, inaccessiblemem: none) uwtable
define noalias noundef ptr @reduce.33(ptr nocapture readonly %0) local_unnamed_addr #0 {
%args_gep = getelementptr inbounds nuw i8, ptr %0, i64 24
%args = load ptr, ptr %args_gep, align 8
%arg0 = load ptr, ptr %args, align 8, !invariant.load !0, !dereferenceable !1, !align !2
%arg1_gep = getelementptr i8, ptr %args, i64 16
%arg1 = load ptr, ptr %arg1_gep, align 8, !invariant.load !0, !dereferenceable !3, !align !2
%arg2_gep = getelementptr i8, ptr %args, i64 32
%arg2 = load ptr, ptr %arg2_gep, align 8, !invariant.load !0, !dereferenceable !2, !align !2
br label %vector.ph
vector.ph: ; preds = %vector.ph, %1
%reduce.33.invar_address.dim.0.07 = phi i64 [ 0, %1 ], [ %invar.inc, %vector.ph ]
%accumulator_0.03 = load float, ptr %arg1, align 64
%2 = insertelement <8 x float> <float poison, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00, float -0.000000e+00>, float %accumulator_0.03, i64 0
%3 = getelementptr inbounds nuw [16 x [16 x float]], ptr %arg0, i64 0, i64 %reduce.33.invar_address.dim.0.07, i64 0
%4 = getelementptr inbounds nuw i8, ptr %3, i64 32
%wide.load10 = load <8 x float>, ptr %4, align 32, !invariant.load !0, !noalias !4
%wide.load = load <8 x float>, ptr %3, align 64, !invariant.load !0, !noalias !4
%5 = fadd reassoc <8 x float> %2, %wide.load
%bin.rdx = fadd reassoc <8 x float> %wide.load10, %5
%6 = tail call reassoc float @llvm.vector.reduce.fadd.v8f32(float -0.000000e+00, <8 x float> %bin.rdx)
%7 = getelementptr inbounds nuw [16 x float], ptr %arg2, i64 0, i64 %reduce.33.invar_address.dim.0.07
store float %6, ptr %7, align 4, !alias.scope !4
%invar.inc = add nuw nsw i64 %reduce.33.invar_address.dim.0.07, 1
%exitcond8 = icmp eq i64 %invar.inc, 16
br i1 %exitcond8, label %return, label %vector.ph
return: ; preds = %vector.ph
ret ptr null
}
; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
declare float @llvm.vector.reduce.fadd.v8f32(float, <8 x float>) #1
attributes #0 = { nofree norecurse nosync nounwind memory(readwrite, inaccessiblemem: none) uwtable "frame-pointer"="all" "prefer-vector-width"="256" }
attributes #1 = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
!0 = !{}
!1 = !{i64 1024}
!2 = !{i64 64}
!3 = !{i64 4}
!4 = !{!5}
!5 = !{!"result slice: {index:12, offset:0, size:64}", !6}
!6 = !{!"XLA host kernel reduce.33 AA domain"}
; ==== module part 03 of 04 ====
; ModuleID = '__compute_module_part_03'
source_filename = "__compute_module"
target datalayout = "e-m:o-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-darwin25.5.0"
; Function Attrs: nofree norecurse nosync nounwind memory(readwrite, inaccessiblemem: none) uwtable
define noalias noundef ptr @broadcast_divide_fusion(ptr nocapture readonly %0) local_unnamed_addr #0 {
%args_gep = getelementptr inbounds nuw i8, ptr %0, i64 24
%args = load ptr, ptr %args_gep, align 8
%arg0 = load ptr, ptr %args, align 8, !invariant.load !0, !dereferenceable !1, !align !2
%arg1_gep = getelementptr i8, ptr %args, i64 16
%arg1 = load ptr, ptr %arg1_gep, align 8, !invariant.load !0, !dereferenceable !2, !align !2
%arg2_gep = getelementptr i8, ptr %args, i64 32
%arg2 = load ptr, ptr %arg2_gep, align 8, !invariant.load !0, !dereferenceable !1, !align !2
br label %vector.ph
vector.ph: ; preds = %vector.ph, %1
%broadcast_divide_fusion.invar_address.dim.0.02 = phi i64 [ 0, %1 ], [ %invar.inc, %vector.ph ]
%2 = getelementptr inbounds nuw [16 x float], ptr %arg1, i64 0, i64 %broadcast_divide_fusion.invar_address.dim.0.02
%3 = load float, ptr %2, align 4, !invariant.load !0, !noalias !3
%broadcast.splatinsert = insertelement <8 x float> poison, float %3, i64 0
%broadcast.splat = shufflevector <8 x float> %broadcast.splatinsert, <8 x float> poison, <8 x i32> zeroinitializer
%4 = getelementptr inbounds nuw [16 x [16 x float]], ptr %arg0, i64 0, i64 %broadcast_divide_fusion.invar_address.dim.0.02, i64 0
%5 = getelementptr inbounds nuw i8, ptr %4, i64 32
%wide.load = load <8 x float>, ptr %4, align 64, !invariant.load !0, !noalias !3
%wide.load4 = load <8 x float>, ptr %5, align 32, !invariant.load !0, !noalias !3
%6 = fdiv <8 x float> %wide.load, %broadcast.splat
%7 = fdiv <8 x float> %wide.load4, %broadcast.splat
%8 = getelementptr inbounds nuw [16 x [16 x float]], ptr %arg2, i64 0, i64 %broadcast_divide_fusion.invar_address.dim.0.02, i64 0
%9 = getelementptr inbounds nuw i8, ptr %8, i64 32
store <8 x float> %6, ptr %8, align 64, !alias.scope !3
store <8 x float> %7, ptr %9, align 32, !alias.scope !3
%invar.inc = add nuw nsw i64 %broadcast_divide_fusion.invar_address.dim.0.02, 1
%exitcond3 = icmp eq i64 %invar.inc, 16
br i1 %exitcond3, label %return, label %vector.ph
return: ; preds = %vector.ph
ret ptr null
}
attributes #0 = { nofree norecurse nosync nounwind memory(readwrite, inaccessiblemem: none) uwtable "frame-pointer"="all" "prefer-vector-width"="256" }
!0 = !{}
!1 = !{i64 1024}
!2 = !{i64 64}
!3 = !{!4}
!4 = !{!"result slice: {index:12, offset:2048, size:1024}", !5}
!5 = !{!"XLA host kernel broadcast_divide_fusion AA domain"}bench/specimen/artifacts/llvm-ir-with-opt.ll
$L__BB0_3: // =>This Inner Loop Header: Depth=1
setp.lt.s64 %p13, %rd92, %rd1;
add.s32 %r225, %r379, 1;
setp.gt.s32 %p14, %r225, 1;
selp.b32 %r379, 0, %r225, %p14;
.loc 1 182 20 // generate_runner_artifacts.py:182:20
cp.async.wait_group 2;
bar.sync 0;
shl.b32 %r226, %r379, 14;
add.s32 %r227, %r93, %r226;
add.s32 %r228, %r227, %r14;
ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%r112, %r113, %r114, %r115}, [%r228];
ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%r124, %r125, %r126, %r127}, [%r228+4096];
ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%r128, %r129, %r130, %r131}, [%r228+8192];
ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%r132, %r133, %r134, %r135}, [%r228+12288];
add.s32 %r229, %r227, %r15;
ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%r136, %r137, %r138, %r139}, [%r229];
ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%r148, %r149, %r150, %r151}, [%r229+4096];
ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%r152, %r153, %r154, %r155}, [%r229+8192];
ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%r156, %r157, %r158, %r159}, [%r229+12288];
add.s32 %r230, %r227, %r16;
ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%r160, %r161, %r162, %r163}, [%r230];
ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%r172, %r173, %r174, %r175}, [%r230+4096];
ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%r176, %r177, %r178, %r179}, [%r230+8192];
ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%r180, %r181, %r182, %r183}, [%r230+12288];
add.s32 %r231, %r227, %r17;
ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%r184, %r185, %r186, %r187}, [%r231];
ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%r196, %r197, %r198, %r199}, [%r231+4096];
ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%r200, %r201, %r202, %r203}, [%r231+8192];
ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%r204, %r205, %r206, %r207}, [%r231+12288];
.loc 1 183 20 // generate_runner_artifacts.py:183:20
shl.b32 %r232, %r379, 13;
add.s32 %r233, %r93, %r232;
add.s32 %r234, %r233, %r18;
ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%r116, %r117, %r140, %r141}, [%r234+32768];
ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%r164, %r165, %r188, %r189}, [%r234+36864];
add.s32 %r235, %r233, %r19;
ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%r118, %r119, %r142, %r143}, [%r235+32768];
ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%r166, %r167, %r190, %r191}, [%r235+36864];
add.s32 %r236, %r233, %r20;
ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%r120, %r121, %r144, %r145}, [%r236+32768];
ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%r168, %r169, %r192, %r193}, [%r236+36864];
add.s32 %r237, %r233, %r21;
ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%r122, %r123, %r146, %r147}, [%r237+32768];
ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%r170, %r171, %r194, %r195}, [%r237+36864];
.loc 1 184 32 // generate_runner_artifacts.py:184:32
// begin inline asm
mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %r381, %r382, %r383, %r384 }, { %r112, %r113, %r114, %r115 }, { %r116, %r117 }, { %r381, %r382, %r383, %r384 };
// end inline asm
// begin inline asm
mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %r385, %r386, %r387, %r388 }, { %r112, %r113, %r114, %r115 }, { %r118, %r119 }, { %r385, %r386, %r387, %r388 };
// end inline asm
// begin inline asm
mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %r389, %r390, %r391, %r392 }, { %r112, %r113, %r114, %r115 }, { %r120, %r121 }, { %r389, %r390, %r391, %r392 };
// end inline asm
// begin inline asm
mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %r393, %r394, %r395, %r396 }, { %r112, %r113, %r114, %r115 }, { %r122, %r123 }, { %r393, %r394, %r395, %r396 };
// end inline asm
// begin inline asm
mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %r397, %r398, %r399, %r400 }, { %r124, %r125, %r126, %r127 }, { %r116, %r117 }, { %r397, %r398, %r399, %r400 };
// end inline asm
// begin inline asm
mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %r401, %r402, %r403, %r404 }, { %r124, %r125, %r126, %r127 }, { %r118, %r119 }, { %r401, %r402, %r403, %r404 };
// end inline asm
// begin inline asm
mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %r405, %r406, %r407, %r408 }, { %r124, %r125, %r126, %r127 }, { %r120, %r121 }, { %r405, %r406, %r407, %r408 };
// end inline asm
// begin inline asm
mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %r409, %r410, %r411, %r412 }, { %r124, %r125, %r126, %r127 }, { %r122, %r123 }, { %r409, %r410, %r411, %r412 };bench/specimen/artifacts/ptx.ptx
Fatbin elf code:
================
arch = sm_80
code version = [1,7]
host = linux
compile_size = 64bit
code for sm_80
Function : _Z17warp_mma_16x16x16P6__halfS0_Pf
.headerflags @"EF_CUDA_SM80 EF_CUDA_PTX_SM(EF_CUDA_SM80)"
/*0000*/ IMAD.MOV.U32 R1, RZ, RZ, c[0x0][0x28] ; /* 0x00000a00ff017624 */
/* 0x000fc400078e00ff */
/*0010*/ S2R R5, SR_LANEID ; /* 0x0000000000057919 */
/* 0x000e220000000000 */
/*0020*/ IMAD.MOV.U32 R3, RZ, RZ, RZ ; /* 0x000000ffff037224 */
/* 0x000fe200078e00ff */
/*0030*/ ULDC.64 UR4, c[0x0][0x118] ; /* 0x0000460000047ab9 */
/* 0x000fe20000000a00 */
/*0040*/ LOP3.LUT R2, R5, 0x3, RZ, 0xc0, !PT ; /* 0x0000000305027812 */
/* 0x001fe400078ec0ff */
/*0050*/ SHF.R.U32.HI R5, RZ, 0x2, R5 ; /* 0x00000002ff057819 */
/* 0x000fca0000011605 */
/*0060*/ IMAD.WIDE.U32 R2, R5, 0x8, R2 ; /* 0x0000000805027825 */
/* 0x000fca00078e0002 */
/*0070*/ LEA R8, P0, R2.reuse, c[0x0][0x160], 0x2 ; /* 0x0000580002087a11 */
/* 0x040fe400078010ff */
/*0080*/ LEA R10, P1, R2.reuse, c[0x0][0x168], 0x2 ; /* 0x00005a00020a7a11 */
/* 0x040fe400078210ff */
/*0090*/ LEA.HI.X R9, R2.reuse, c[0x0][0x164], R3.reuse, 0x2, P0 ; /* 0x0000590002097a11 */
/* 0x140fe400000f1403 */
/*00a0*/ LEA.HI.X R11, R2, c[0x0][0x16c], R3, 0x2, P1 ; /* 0x00005b00020b7a11 */
/* 0x000fc600008f1403 */
/*00b0*/ LDG.E R4, [R8.64] ; /* 0x0000000408047981 */
/* 0x000ea8000c1e1900 */
/*00c0*/ LDG.E R5, [R8.64+0x100] ; /* 0x0001000408057981 */
/* 0x000ea8000c1e1900 */
/*00d0*/ LDG.E R6, [R8.64+0x10] ; /* 0x0000100408067981 */
/* 0x000ea8000c1e1900 */
/*00e0*/ LDG.E R7, [R8.64+0x110] ; /* 0x0001100408077981 */
/* 0x000ea8000c1e1900 */
/*00f0*/ LDG.E R12, [R10.64] ; /* 0x000000040a0c7981 */
/* 0x000ea8000c1e1900 */
/*0100*/ LDG.E R13, [R10.64+0x10] ; /* 0x000010040a0d7981 */
/* 0x000ea8000c1e1900 */
/*0110*/ LDG.E R14, [R10.64+0x100] ; /* 0x000100040a0e7981 */
/* 0x000ee8000c1e1900 */
/*0120*/ LDG.E R15, [R10.64+0x110] ; /* 0x000110040a0f7981 */
/* 0x000ee2000c1e1900 */
/*0130*/ HMMA.16816.F32 R16, R4.reuse, R12, RZ ; /* 0x0000000c0410723c */
/* 0x044f7000000018ff */
/*0140*/ HMMA.16816.F32 R12, R4, R14, RZ ; /* 0x0000000e040c723c */
/* 0x008b6e00000018ff */
/*0150*/ LEA R4, P0, R2, c[0x0][0x170], 0x3 ; /* 0x00005c0002047a11 */
/* 0x020fc800078018ff */
/*0160*/ LEA.HI.X R5, R2, c[0x0][0x174], R3, 0x3, P0 ; /* 0x00005d0002057a11 */
/* 0x000fca00000f1c03 */
/*0170*/ STG.E.64 [R4.64], R16 ; /* 0x0000001004007986 */
/* 0x000fe8000c101b04 */
/*0180*/ STG.E.64 [R4.64+0x200], R18 ; /* 0x0002001204007986 */
/* 0x000fe8000c101b04 */
/*0190*/ STG.E.64 [R4.64+0x20], R12 ; /* 0x0000200c04007986 */
/* 0x000fe8000c101b04 */
/*01a0*/ STG.E.64 [R4.64+0x220], R14 ; /* 0x0002200e04007986 */
/* 0x000fe2000c101b04 */
/*01b0*/ EXIT ; /* 0x000000000000794d */
/* 0x000fea0003800000 */
/*01c0*/ BRA 0x1c0; /* 0xfffffff000007947 */
/* 0x000fc0000383ffff */bench/specimen/artifacts/sass.txt
j and k step too
captured by python3 bench/specimen/capture.py · every level below the source is machine output
pending
00
00
The levels this machine cannot reach
A CPU has no Triton stack and no torch_xla device, so four floors of the map stay empty here. Each one waits on a specific capture rather than on a drawing of what the dump would probably look like.
| level | lane | artifact | what will produce it |
|---|
provenance
What produced these files
- command
- python3 bench/specimen/capture.py, from the repo root
- program
- bench/specimen/specimen.py, seed 0, the block sums to -269.723755 on this run
- versions
- python 3.12.0 · jax 0.4.38 · jaxlib 0.4.38
- machine
- macOS-26.5.2-x86_64-i386-64bit · cpu backend · TFRT_CPU_0
- edits
- HLO instruction metadata keeps op_name and drops source_file and source_line, because those name the machine that ran the compile rather than the program. The LLVM IR and the StableHLO module carry no paths and are written verbatim. Nothing else is edited.
- files
- bench/specimen/artifacts/source.py (40 lines) · bench/specimen/artifacts/jaxpr.txt (42 lines) · bench/specimen/artifacts/torch-trace.txt (67 lines) · bench/specimen/artifacts/stablehlo.mlir (29 lines) · bench/specimen/artifacts/hlo-before-optimizations.txt (48 lines) · bench/specimen/artifacts/hlo-after-optimizations.txt (61 lines) · bench/specimen/artifacts/triton-ttir.mlir (68 lines) · bench/specimen/artifacts/triton-ttgir.mlir (98 lines) · bench/specimen/artifacts/llvm-ir-no-opt.ll (479 lines) · bench/specimen/artifacts/llvm-ir-with-opt.ll (282 lines) · bench/specimen/artifacts/ptx.ptx (69 lines) · bench/specimen/artifacts/sass.txt (68 lines)
The capture script clears its output directory before it writes, so running it again on
the same versions leaves the committed files unchanged. LAB·L1 walks the same descent
cell by cell in a notebook.