Model
Self-attention with dynamic batch and sequence dimensions (tensor<?x?x64xf32>), 3 weight projections (Q/K/V), scaled dot-product attention with softmax.
Pipeline
hip-mlir-opt attention.onnx.mlir \
--allow-unregistered-dialect \
--onnx-to-hip-pipeline='externalize-min-num-elements=256 externalize-output-dir=/tmp' \
--mlir-elide-elementsattrs-if-larger=4 \
--mlir-print-ir-after-all
21 passes executed in the onnx-to-hip-pipeline. Complete IR after each pass is shown below.
Pass 1: hip-add-context-arg
Inserts !hip.context as the first argument to @main_graph. The ONNX ops are unchanged.
module attributes {llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128", llvm.target_triple = "x86_64-unknown-linux-gnu", "onnx-mlir.symbol-postfix" = "attention_dynamic"} {
func.func @main_graph(%arg0: !hip.context, %arg1: tensor<?x?x64xf32> {onnx.dim_params = "0:batch,1:seq", onnx.name = "X"}) -> (tensor<?x?x64xf32> {onnx.dim_params = "0:batch,1:seq", onnx.name = "out"}) {
%0 = "onnx.Constant"() {value = dense_resource<__elided__> : tensor<64x64xf32>} : () -> tensor<64x64xf32>
%1 = "onnx.Constant"() {value = dense_resource<__elided__> : tensor<64x64xf32>} : () -> tensor<64x64xf32>
%2 = "onnx.Constant"() {value = dense_resource<__elided__> : tensor<64x64xf32>} : () -> tensor<64x64xf32>
%3 = "onnx.Constant"() {value = dense<1.250000e-01> : tensor<f32>} : () -> tensor<f32>
%4 = "onnx.MatMul"(%arg1, %0) {onnx_node_name = "q_proj"} : (tensor<?x?x64xf32>, tensor<64x64xf32>) -> tensor<?x?x64xf32>
%5 = "onnx.MatMul"(%arg1, %1) {onnx_node_name = "k_proj"} : (tensor<?x?x64xf32>, tensor<64x64xf32>) -> tensor<?x?x64xf32>
%6 = "onnx.MatMul"(%arg1, %2) {onnx_node_name = "v_proj"} : (tensor<?x?x64xf32>, tensor<64x64xf32>) -> tensor<?x?x64xf32>
%7 = "onnx.Transpose"(%5) <{perm = [0, 2, 1]}> {onnx_node_name = "k_transpose"} : (tensor<?x?x64xf32>) -> tensor<?x64x?xf32>
%8 = "onnx.MatMul"(%4, %7) {onnx_node_name = "qk_matmul"} : (tensor<?x?x64xf32>, tensor<?x64x?xf32>) -> tensor<?x?x?xf32>
%9 = "onnx.Mul"(%8, %3) {onnx_node_name = "score_scale"} : (tensor<?x?x?xf32>, tensor<f32>) -> tensor<?x?x?xf32>
%10 = "onnx.Softmax"(%9) <{axis = -1 : si64}> {onnx_node_name = "softmax"} : (tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
%11 = "onnx.MatMul"(%10, %6) {onnx_node_name = "attn_v"} : (tensor<?x?x?xf32>, tensor<?x?x64xf32>) -> tensor<?x?x64xf32>
return %11 : tensor<?x?x64xf32>
}
"onnx.EntryPoint"() <{func = @main_graph}> : () -> ()
}
Pass 2: convert-onnx-to-hip
Converts ONNX ops to HIP dialect ops (tensor mode). Large constants (>= 256 elements) are externalized to /tmp/model.constants.bin and replaced with memref.global + hip.external_data attributes. Small constants are kept inline as arith.constant.
module attributes {hip.constants_file = "model.constants.bin", llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128", llvm.target_triple = "x86_64-unknown-linux-gnu", "onnx-mlir.symbol-postfix" = "attention_dynamic"} {
memref.global "private" @hip_ext_constant_2 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 32768 : i64, size = 16384 : i64}}
memref.global "private" @hip_ext_constant_1 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 16384 : i64, size = 16384 : i64}}
memref.global "private" @hip_ext_constant_0 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 0 : i64, size = 16384 : i64}}
func.func @main_graph(%arg0: !hip.context, %arg1: tensor<?x?x64xf32> {onnx.dim_params = "0:batch,1:seq", onnx.name = "X"}) -> (tensor<?x?x64xf32> {onnx.dim_params = "0:batch,1:seq", onnx.name = "out"}) {
%c2 = arith.constant 2 : index
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%cst = arith.constant dense<1.250000e-01> : tensor<f32>
%0 = memref.get_global @hip_ext_constant_0 : memref<64x64xf32>
%1 = bufferization.to_tensor %0 restrict : memref<64x64xf32> to tensor<64x64xf32>
%2 = memref.get_global @hip_ext_constant_1 : memref<64x64xf32>
%3 = bufferization.to_tensor %2 restrict : memref<64x64xf32> to tensor<64x64xf32>
%4 = memref.get_global @hip_ext_constant_2 : memref<64x64xf32>
%5 = bufferization.to_tensor %4 restrict : memref<64x64xf32> to tensor<64x64xf32>
%dim = tensor.dim %arg1, %c0 : tensor<?x?x64xf32>
%dim_0 = tensor.dim %arg1, %c1 : tensor<?x?x64xf32>
%6 = tensor.empty(%dim, %dim_0) : tensor<?x?x64xf32>
%7 = hip.hipblaslt.matmul(%arg0) ins(%arg1, %1 : tensor<?x?x64xf32>, tensor<64x64xf32>) outs(%6 : tensor<?x?x64xf32>) -> tensor<?x?x64xf32>
%dim_1 = tensor.dim %arg1, %c0 : tensor<?x?x64xf32>
%dim_2 = tensor.dim %arg1, %c1 : tensor<?x?x64xf32>
%8 = tensor.empty(%dim_1, %dim_2) : tensor<?x?x64xf32>
%9 = hip.hipblaslt.matmul(%arg0) ins(%arg1, %3 : tensor<?x?x64xf32>, tensor<64x64xf32>) outs(%8 : tensor<?x?x64xf32>) -> tensor<?x?x64xf32>
%dim_3 = tensor.dim %arg1, %c0 : tensor<?x?x64xf32>
%dim_4 = tensor.dim %arg1, %c1 : tensor<?x?x64xf32>
%10 = tensor.empty(%dim_3, %dim_4) : tensor<?x?x64xf32>
%11 = hip.hipblaslt.matmul(%arg0) ins(%arg1, %5 : tensor<?x?x64xf32>, tensor<64x64xf32>) outs(%10 : tensor<?x?x64xf32>) -> tensor<?x?x64xf32>
%dim_5 = tensor.dim %9, %c0 : tensor<?x?x64xf32>
%dim_6 = tensor.dim %9, %c1 : tensor<?x?x64xf32>
%12 = tensor.empty(%dim_5, %dim_6) : tensor<?x64x?xf32>
%13 = hip.transpose(%arg0, %c1, %c2) ins(%9 : tensor<?x?x64xf32>) outs(%12 : tensor<?x64x?xf32>) -> tensor<?x64x?xf32>
%dim_7 = tensor.dim %7, %c0 : tensor<?x?x64xf32>
%dim_8 = tensor.dim %7, %c1 : tensor<?x?x64xf32>
%dim_9 = tensor.dim %13, %c2 : tensor<?x64x?xf32>
%14 = tensor.empty(%dim_7, %dim_8, %dim_9) : tensor<?x?x?xf32>
%15 = hip.hipblaslt.matmul(%arg0) ins(%7, %13 : tensor<?x?x64xf32>, tensor<?x64x?xf32>) outs(%14 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
%dim_10 = tensor.dim %15, %c0 : tensor<?x?x?xf32>
%dim_11 = tensor.dim %15, %c1 : tensor<?x?x?xf32>
%dim_12 = tensor.dim %15, %c2 : tensor<?x?x?xf32>
%16 = tensor.empty(%dim_10, %dim_11, %dim_12) : tensor<?x?x?xf32>
%17 = hip.miopen.mul(%arg0) ins(%15, %cst : tensor<?x?x?xf32>, tensor<f32>) outs(%16 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
%dim_13 = tensor.dim %17, %c0 : tensor<?x?x?xf32>
%dim_14 = tensor.dim %17, %c1 : tensor<?x?x?xf32>
%dim_15 = tensor.dim %17, %c2 : tensor<?x?x?xf32>
%18 = tensor.empty(%dim_13, %dim_14, %dim_15) : tensor<?x?x?xf32>
%19 = hip.miopen.softmax(%arg0) ins(%17 : tensor<?x?x?xf32>) outs(%18 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
%dim_16 = tensor.dim %19, %c0 : tensor<?x?x?xf32>
%dim_17 = tensor.dim %19, %c1 : tensor<?x?x?xf32>
%20 = tensor.empty(%dim_16, %dim_17) : tensor<?x?x64xf32>
%21 = hip.hipblaslt.matmul(%arg0) ins(%19, %11 : tensor<?x?x?xf32>, tensor<?x?x64xf32>) outs(%20 : tensor<?x?x64xf32>) -> tensor<?x?x64xf32>
return %21 : tensor<?x?x64xf32>
}
}
Pass 3: one-shot-bufferize
Converts tensor operations to memref operations. tensor.empty becomes memref.alloc, tensor.dim becomes memref.dim. Function signature changes from tensor returns to memref returns. 8 memref.alloc calls are introduced.
module attributes {hip.constants_file = "model.constants.bin", llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128", llvm.target_triple = "x86_64-unknown-linux-gnu", "onnx-mlir.symbol-postfix" = "attention_dynamic"} {
memref.global "private" constant @__constant_xf32 : memref<f32> = dense<1.250000e-01> {alignment = 64 : i64}
memref.global "private" @hip_ext_constant_2 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 32768 : i64, size = 16384 : i64}}
memref.global "private" @hip_ext_constant_1 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 16384 : i64, size = 16384 : i64}}
memref.global "private" @hip_ext_constant_0 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 0 : i64, size = 16384 : i64}}
func.func @main_graph(%arg0: !hip.context, %arg1: memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>> {onnx.dim_params = "0:batch,1:seq", onnx.name = "X"}) -> (memref<?x?x64xf32> {onnx.dim_params = "0:batch,1:seq", onnx.name = "out"}) {
%c2 = arith.constant 2 : index
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%0 = memref.get_global @__constant_xf32 : memref<f32>
%1 = memref.get_global @hip_ext_constant_0 : memref<64x64xf32>
%2 = memref.get_global @hip_ext_constant_1 : memref<64x64xf32>
%3 = memref.get_global @hip_ext_constant_2 : memref<64x64xf32>
%dim = memref.dim %arg1, %c0 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%dim_0 = memref.dim %arg1, %c1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%alloc = memref.alloc(%dim, %dim_0) {alignment = 64 : i64} : memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%alloc : memref<?x?x64xf32>)
%dim_1 = memref.dim %arg1, %c0 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%dim_2 = memref.dim %arg1, %c1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%alloc_3 = memref.alloc(%dim_1, %dim_2) {alignment = 64 : i64} : memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %2 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%alloc_3 : memref<?x?x64xf32>)
%dim_4 = memref.dim %arg1, %c0 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%dim_5 = memref.dim %arg1, %c1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%alloc_6 = memref.alloc(%dim_4, %dim_5) {alignment = 64 : i64} : memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %3 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%alloc_6 : memref<?x?x64xf32>)
%dim_7 = memref.dim %alloc_3, %c0 : memref<?x?x64xf32>
%dim_8 = memref.dim %alloc_3, %c1 : memref<?x?x64xf32>
%alloc_9 = memref.alloc(%dim_7, %dim_8) {alignment = 64 : i64} : memref<?x64x?xf32>
hip.transpose(%arg0, %c1, %c2) ins(%alloc_3 : memref<?x?x64xf32>) outs(%alloc_9 : memref<?x64x?xf32>)
%dim_10 = memref.dim %alloc, %c0 : memref<?x?x64xf32>
%dim_11 = memref.dim %alloc, %c1 : memref<?x?x64xf32>
%dim_12 = memref.dim %alloc_9, %c2 : memref<?x64x?xf32>
%alloc_13 = memref.alloc(%dim_10, %dim_11, %dim_12) {alignment = 64 : i64} : memref<?x?x?xf32>
hip.hipblaslt.matmul(%arg0) ins(%alloc, %alloc_9 : memref<?x?x64xf32>, memref<?x64x?xf32>) outs(%alloc_13 : memref<?x?x?xf32>)
%dim_14 = memref.dim %alloc_13, %c0 : memref<?x?x?xf32>
%dim_15 = memref.dim %alloc_13, %c1 : memref<?x?x?xf32>
%dim_16 = memref.dim %alloc_13, %c2 : memref<?x?x?xf32>
%alloc_17 = memref.alloc(%dim_14, %dim_15, %dim_16) {alignment = 64 : i64} : memref<?x?x?xf32>
hip.miopen.mul(%arg0) ins(%alloc_13, %0 : memref<?x?x?xf32>, memref<f32>) outs(%alloc_17 : memref<?x?x?xf32>)
%dim_18 = memref.dim %alloc_17, %c0 : memref<?x?x?xf32>
%dim_19 = memref.dim %alloc_17, %c1 : memref<?x?x?xf32>
%dim_20 = memref.dim %alloc_17, %c2 : memref<?x?x?xf32>
%alloc_21 = memref.alloc(%dim_18, %dim_19, %dim_20) {alignment = 64 : i64} : memref<?x?x?xf32>
hip.miopen.softmax(%arg0) ins(%alloc_17 : memref<?x?x?xf32>) outs(%alloc_21 : memref<?x?x?xf32>)
%dim_22 = memref.dim %alloc_21, %c0 : memref<?x?x?xf32>
%dim_23 = memref.dim %alloc_21, %c1 : memref<?x?x?xf32>
%alloc_24 = memref.alloc(%dim_22, %dim_23) {alignment = 64 : i64} : memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%alloc_21, %alloc_6 : memref<?x?x?xf32>, memref<?x?x64xf32>) outs(%alloc_24 : memref<?x?x64xf32>)
%cast = memref.cast %alloc_24 : memref<?x?x64xf32> to memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
return %alloc_24 : memref<?x?x64xf32>
}
}
Pass 4: buffer-results-to-out-params
Converts the memref return value to an output parameter (%arg2). The function now returns void. The final matmul writes directly to %arg2 instead of a locally-allocated buffer.
module attributes {hip.constants_file = "model.constants.bin", llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128", llvm.target_triple = "x86_64-unknown-linux-gnu", "onnx-mlir.symbol-postfix" = "attention_dynamic"} {
memref.global "private" constant @__constant_xf32 : memref<f32> = dense<1.250000e-01> {alignment = 64 : i64}
memref.global "private" @hip_ext_constant_2 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 32768 : i64, size = 16384 : i64}}
memref.global "private" @hip_ext_constant_1 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 16384 : i64, size = 16384 : i64}}
memref.global "private" @hip_ext_constant_0 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 0 : i64, size = 16384 : i64}}
func.func @main_graph(%arg0: !hip.context, %arg1: memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>> {onnx.dim_params = "0:batch,1:seq", onnx.name = "X"}, %arg2: memref<?x?x64xf32> {bufferize.result, onnx.dim_params = "0:batch,1:seq", onnx.name = "out"}) {
%c2 = arith.constant 2 : index
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%0 = memref.get_global @__constant_xf32 : memref<f32>
%1 = memref.get_global @hip_ext_constant_0 : memref<64x64xf32>
%2 = memref.get_global @hip_ext_constant_1 : memref<64x64xf32>
%3 = memref.get_global @hip_ext_constant_2 : memref<64x64xf32>
%dim = memref.dim %arg1, %c0 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%dim_0 = memref.dim %arg1, %c1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%alloc = memref.alloc(%dim, %dim_0) {alignment = 64 : i64} : memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%alloc : memref<?x?x64xf32>)
%dim_1 = memref.dim %arg1, %c0 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%dim_2 = memref.dim %arg1, %c1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%alloc_3 = memref.alloc(%dim_1, %dim_2) {alignment = 64 : i64} : memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %2 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%alloc_3 : memref<?x?x64xf32>)
%dim_4 = memref.dim %arg1, %c0 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%dim_5 = memref.dim %arg1, %c1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%alloc_6 = memref.alloc(%dim_4, %dim_5) {alignment = 64 : i64} : memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %3 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%alloc_6 : memref<?x?x64xf32>)
%dim_7 = memref.dim %alloc_3, %c0 : memref<?x?x64xf32>
%dim_8 = memref.dim %alloc_3, %c1 : memref<?x?x64xf32>
%alloc_9 = memref.alloc(%dim_7, %dim_8) {alignment = 64 : i64} : memref<?x64x?xf32>
hip.transpose(%arg0, %c1, %c2) ins(%alloc_3 : memref<?x?x64xf32>) outs(%alloc_9 : memref<?x64x?xf32>)
%dim_10 = memref.dim %alloc, %c0 : memref<?x?x64xf32>
%dim_11 = memref.dim %alloc, %c1 : memref<?x?x64xf32>
%dim_12 = memref.dim %alloc_9, %c2 : memref<?x64x?xf32>
%alloc_13 = memref.alloc(%dim_10, %dim_11, %dim_12) {alignment = 64 : i64} : memref<?x?x?xf32>
hip.hipblaslt.matmul(%arg0) ins(%alloc, %alloc_9 : memref<?x?x64xf32>, memref<?x64x?xf32>) outs(%alloc_13 : memref<?x?x?xf32>)
%dim_14 = memref.dim %alloc_13, %c0 : memref<?x?x?xf32>
%dim_15 = memref.dim %alloc_13, %c1 : memref<?x?x?xf32>
%dim_16 = memref.dim %alloc_13, %c2 : memref<?x?x?xf32>
%alloc_17 = memref.alloc(%dim_14, %dim_15, %dim_16) {alignment = 64 : i64} : memref<?x?x?xf32>
hip.miopen.mul(%arg0) ins(%alloc_13, %0 : memref<?x?x?xf32>, memref<f32>) outs(%alloc_17 : memref<?x?x?xf32>)
%dim_18 = memref.dim %alloc_17, %c0 : memref<?x?x?xf32>
%dim_19 = memref.dim %alloc_17, %c1 : memref<?x?x?xf32>
%dim_20 = memref.dim %alloc_17, %c2 : memref<?x?x?xf32>
%alloc_21 = memref.alloc(%dim_18, %dim_19, %dim_20) {alignment = 64 : i64} : memref<?x?x?xf32>
hip.miopen.softmax(%arg0) ins(%alloc_17 : memref<?x?x?xf32>) outs(%alloc_21 : memref<?x?x?xf32>)
%dim_22 = memref.dim %alloc_21, %c0 : memref<?x?x?xf32>
%dim_23 = memref.dim %alloc_21, %c1 : memref<?x?x?xf32>
hip.hipblaslt.matmul(%arg0) ins(%alloc_21, %alloc_6 : memref<?x?x?xf32>, memref<?x?x64xf32>) outs(%arg2 : memref<?x?x64xf32>)
%cast = memref.cast %arg2 : memref<?x?x64xf32> to memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
return
}
}
Pass 5: expand-realloc
No-op for this model (no realloc operations present). IR is identical to Pass 4.
Pass 6: canonicalize
Folds redundant memref.dim operations (dim_1/dim_2 CSE'd with dim/dim_0), removes unused memref.cast, simplifies dim queries. Reduces from 8 to 7 memref.alloc (the output alloc was already replaced by %arg2).
module attributes {hip.constants_file = "model.constants.bin", llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128", llvm.target_triple = "x86_64-unknown-linux-gnu", "onnx-mlir.symbol-postfix" = "attention_dynamic"} {
memref.global "private" constant @__constant_xf32 : memref<f32> = dense<1.250000e-01> {alignment = 64 : i64}
memref.global "private" @hip_ext_constant_2 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 32768 : i64, size = 16384 : i64}}
memref.global "private" @hip_ext_constant_1 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 16384 : i64, size = 16384 : i64}}
memref.global "private" @hip_ext_constant_0 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 0 : i64, size = 16384 : i64}}
func.func @main_graph(%arg0: !hip.context, %arg1: memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>> {onnx.dim_params = "0:batch,1:seq", onnx.name = "X"}, %arg2: memref<?x?x64xf32> {bufferize.result, onnx.dim_params = "0:batch,1:seq", onnx.name = "out"}) {
%c2 = arith.constant 2 : index
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%0 = memref.get_global @__constant_xf32 : memref<f32>
%1 = memref.get_global @hip_ext_constant_0 : memref<64x64xf32>
%2 = memref.get_global @hip_ext_constant_1 : memref<64x64xf32>
%3 = memref.get_global @hip_ext_constant_2 : memref<64x64xf32>
%dim = memref.dim %arg1, %c0 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%dim_0 = memref.dim %arg1, %c1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%alloc = memref.alloc(%dim, %dim_0) {alignment = 64 : i64} : memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%alloc : memref<?x?x64xf32>)
%dim_1 = memref.dim %arg1, %c0 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%dim_2 = memref.dim %arg1, %c1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%alloc_3 = memref.alloc(%dim_1, %dim_2) {alignment = 64 : i64} : memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %2 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%alloc_3 : memref<?x?x64xf32>)
%dim_4 = memref.dim %arg1, %c0 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%dim_5 = memref.dim %arg1, %c1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%alloc_6 = memref.alloc(%dim_4, %dim_5) {alignment = 64 : i64} : memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %3 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%alloc_6 : memref<?x?x64xf32>)
%alloc_7 = memref.alloc(%dim_1, %dim_2) {alignment = 64 : i64} : memref<?x64x?xf32>
hip.transpose(%arg0, %c1, %c2) ins(%alloc_3 : memref<?x?x64xf32>) outs(%alloc_7 : memref<?x64x?xf32>)
%alloc_8 = memref.alloc(%dim, %dim_0, %dim_2) {alignment = 64 : i64} : memref<?x?x?xf32>
hip.hipblaslt.matmul(%arg0) ins(%alloc, %alloc_7 : memref<?x?x64xf32>, memref<?x64x?xf32>) outs(%alloc_8 : memref<?x?x?xf32>)
%alloc_9 = memref.alloc(%dim, %dim_0, %dim_2) {alignment = 64 : i64} : memref<?x?x?xf32>
hip.miopen.mul(%arg0) ins(%alloc_8, %0 : memref<?x?x?xf32>, memref<f32>) outs(%alloc_9 : memref<?x?x?xf32>)
%alloc_10 = memref.alloc(%dim, %dim_0, %dim_2) {alignment = 64 : i64} : memref<?x?x?xf32>
hip.miopen.softmax(%arg0) ins(%alloc_9 : memref<?x?x?xf32>) outs(%alloc_10 : memref<?x?x?xf32>)
hip.hipblaslt.matmul(%arg0) ins(%alloc_10, %alloc_6 : memref<?x?x?xf32>, memref<?x?x64xf32>) outs(%arg2 : memref<?x?x64xf32>)
return
}
}
Observation: No memref.dealloc yet -- the buffer deallocation pass hasn't run.
Passes 7-14: Buffer deallocation and cleanup
Passes 7-14 are the standard MLIR buffer deallocation pipeline: ownership-based-buffer-deallocation -> canonicalize -> buffer-deallocation-simplification -> bufferization-lower-deallocations -> cse -> canonicalize -> cse -> canonicalize. The final result after pass 14 is the clean form with 7 memref.alloc / 7 memref.dealloc pairs:
module attributes {hip.constants_file = "model.constants.bin", llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128", llvm.target_triple = "x86_64-unknown-linux-gnu", "onnx-mlir.symbol-postfix" = "attention_dynamic"} {
memref.global "private" constant @__constant_xf32 : memref<f32> = dense<1.250000e-01> {alignment = 64 : i64}
memref.global "private" @hip_ext_constant_2 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 32768 : i64, size = 16384 : i64}}
memref.global "private" @hip_ext_constant_1 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 16384 : i64, size = 16384 : i64}}
memref.global "private" @hip_ext_constant_0 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 0 : i64, size = 16384 : i64}}
func.func @main_graph(%arg0: !hip.context, %arg1: memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>> {onnx.dim_params = "0:batch,1:seq", onnx.name = "X"}, %arg2: memref<?x?x64xf32> {bufferize.result, onnx.dim_params = "0:batch,1:seq", onnx.name = "out"}) {
%c2 = arith.constant 2 : index
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%0 = memref.get_global @__constant_xf32 : memref<f32>
%1 = memref.get_global @hip_ext_constant_0 : memref<64x64xf32>
%2 = memref.get_global @hip_ext_constant_1 : memref<64x64xf32>
%3 = memref.get_global @hip_ext_constant_2 : memref<64x64xf32>
%dim = memref.dim %arg1, %c0 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%dim_0 = memref.dim %arg1, %c1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%alloc = memref.alloc(%dim, %dim_0) {alignment = 64 : i64} : memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%alloc : memref<?x?x64xf32>)
%alloc_1 = memref.alloc(%dim, %dim_0) {alignment = 64 : i64} : memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %2 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%alloc_1 : memref<?x?x64xf32>)
%alloc_2 = memref.alloc(%dim, %dim_0) {alignment = 64 : i64} : memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %3 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%alloc_2 : memref<?x?x64xf32>)
%alloc_3 = memref.alloc(%dim, %dim_0) {alignment = 64 : i64} : memref<?x64x?xf32>
hip.transpose(%arg0, %c1, %c2) ins(%alloc_1 : memref<?x?x64xf32>) outs(%alloc_3 : memref<?x64x?xf32>)
%alloc_4 = memref.alloc(%dim, %dim_0, %dim_0) {alignment = 64 : i64} : memref<?x?x?xf32>
hip.hipblaslt.matmul(%arg0) ins(%alloc, %alloc_3 : memref<?x?x64xf32>, memref<?x64x?xf32>) outs(%alloc_4 : memref<?x?x?xf32>)
%alloc_5 = memref.alloc(%dim, %dim_0, %dim_0) {alignment = 64 : i64} : memref<?x?x?xf32>
hip.miopen.mul(%arg0) ins(%alloc_4, %0 : memref<?x?x?xf32>, memref<f32>) outs(%alloc_5 : memref<?x?x?xf32>)
%alloc_6 = memref.alloc(%dim, %dim_0, %dim_0) {alignment = 64 : i64} : memref<?x?x?xf32>
hip.miopen.softmax(%arg0) ins(%alloc_5 : memref<?x?x?xf32>) outs(%alloc_6 : memref<?x?x?xf32>)
hip.hipblaslt.matmul(%arg0) ins(%alloc_6, %alloc_2 : memref<?x?x?xf32>, memref<?x?x64xf32>) outs(%arg2 : memref<?x?x64xf32>)
memref.dealloc %alloc : memref<?x?x64xf32>
memref.dealloc %alloc_1 : memref<?x?x64xf32>
memref.dealloc %alloc_2 : memref<?x?x64xf32>
memref.dealloc %alloc_3 : memref<?x64x?xf32>
memref.dealloc %alloc_4 : memref<?x?x?xf32>
memref.dealloc %alloc_5 : memref<?x?x?xf32>
memref.dealloc %alloc_6 : memref<?x?x?xf32>
return
}
}
Observation: 7 separate dynamic memref.alloc / memref.dealloc pairs. This is the input to the memory optimization passes.
Pass 15: hip-optimize-memrefs (liveness analysis)
Analyzes buffer liveness and reuses dead buffers. alloc_6 (softmax output) is eliminated -- softmax now writes directly to alloc_4 (which is dead after the QK^T matmul result is consumed by mul). 7 allocs reduced to 6.
func.func @main_graph(%arg0: !hip.context, %arg1: memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>> {onnx.dim_params = "0:batch,1:seq", onnx.name = "X"}, %arg2: memref<?x?x64xf32> {bufferize.result, onnx.dim_params = "0:batch,1:seq", onnx.name = "out"}) {
%c2 = arith.constant 2 : index
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%0 = memref.get_global @__constant_xf32 : memref<f32>
%1 = memref.get_global @hip_ext_constant_0 : memref<64x64xf32>
%2 = memref.get_global @hip_ext_constant_1 : memref<64x64xf32>
%3 = memref.get_global @hip_ext_constant_2 : memref<64x64xf32>
%dim = memref.dim %arg1, %c0 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%dim_0 = memref.dim %arg1, %c1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%alloc = memref.alloc(%dim, %dim_0) {alignment = 64 : i64} : memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%alloc : memref<?x?x64xf32>)
%alloc_1 = memref.alloc(%dim, %dim_0) {alignment = 64 : i64} : memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %2 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%alloc_1 : memref<?x?x64xf32>)
%alloc_2 = memref.alloc(%dim, %dim_0) {alignment = 64 : i64} : memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %3 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%alloc_2 : memref<?x?x64xf32>)
%alloc_3 = memref.alloc(%dim, %dim_0) {alignment = 64 : i64} : memref<?x64x?xf32>
hip.transpose(%arg0, %c1, %c2) ins(%alloc_1 : memref<?x?x64xf32>) outs(%alloc_3 : memref<?x64x?xf32>)
%alloc_4 = memref.alloc(%dim, %dim_0, %dim_0) {alignment = 64 : i64} : memref<?x?x?xf32>
hip.hipblaslt.matmul(%arg0) ins(%alloc, %alloc_3 : memref<?x?x64xf32>, memref<?x64x?xf32>) outs(%alloc_4 : memref<?x?x?xf32>)
%alloc_5 = memref.alloc(%dim, %dim_0, %dim_0) {alignment = 64 : i64} : memref<?x?x?xf32>
hip.miopen.mul(%arg0) ins(%alloc_4, %0 : memref<?x?x?xf32>, memref<f32>) outs(%alloc_5 : memref<?x?x?xf32>)
hip.miopen.softmax(%arg0) ins(%alloc_5 : memref<?x?x?xf32>) outs(%alloc_4 : memref<?x?x?xf32>)
hip.hipblaslt.matmul(%arg0) ins(%alloc_4, %alloc_2 : memref<?x?x?xf32>, memref<?x?x64xf32>) outs(%arg2 : memref<?x?x64xf32>)
memref.dealloc %alloc : memref<?x?x64xf32>
memref.dealloc %alloc_1 : memref<?x?x64xf32>
memref.dealloc %alloc_2 : memref<?x?x64xf32>
memref.dealloc %alloc_3 : memref<?x64x?xf32>
memref.dealloc %alloc_4 : memref<?x?x?xf32>
memref.dealloc %alloc_5 : memref<?x?x?xf32>
return
}
Observation: alloc_6 eliminated. Softmax output reuses alloc_4 (outs(%alloc_4 : memref<?x?x?xf32>)).
Pass 16: hip-pool-allocs (memory pooling)
All 6 dynamic allocations consolidated into a single memref<?xi8> pool. Each buffer becomes a memref.view at a 256-byte aligned offset. Pool size is computed dynamically at runtime.
func.func @main_graph(%arg0: !hip.context, %arg1: memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>> {onnx.dim_params = "0:batch,1:seq", onnx.name = "X"}, %arg2: memref<?x?x64xf32> {bufferize.result, onnx.dim_params = "0:batch,1:seq", onnx.name = "out"}) attributes {hipdnn.buffer_offsets = [0, -1, -1, -1, -1, -1]} {
%c2 = arith.constant 2 : index
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%0 = memref.get_global @__constant_xf32 : memref<f32>
%1 = memref.get_global @hip_ext_constant_0 : memref<64x64xf32>
%2 = memref.get_global @hip_ext_constant_1 : memref<64x64xf32>
%3 = memref.get_global @hip_ext_constant_2 : memref<64x64xf32>
%dim = memref.dim %arg1, %c0 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%dim_0 = memref.dim %arg1, %c1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%c256 = arith.constant 256 : index
%4 = arith.muli %c256, %dim : index
%5 = arith.muli %4, %dim_0 : index
%c4 = arith.constant 4 : index
%6 = arith.muli %c4, %dim : index
%7 = arith.muli %6, %dim_0 : index
%8 = arith.muli %7, %dim_0 : index
%c0_1 = arith.constant 0 : index
%c256_2 = arith.constant 256 : index
%c255 = arith.constant 255 : index
%9 = arith.addi %5, %c255 : index
%10 = arith.divui %9, %c256_2 : index
%11 = arith.muli %10, %c256_2 : index
%c4_3 = arith.constant 4 : index
%12 = arith.muli %11, %c4_3 : index
%13 = arith.addi %c0_1, %12 : index
%c255_4 = arith.constant 255 : index
%14 = arith.addi %8, %c255_4 : index
%15 = arith.divui %14, %c256_2 : index
%16 = arith.muli %15, %c256_2 : index
%c2_5 = arith.constant 2 : index
%17 = arith.muli %16, %c2_5 : index
%18 = arith.addi %13, %17 : index
%alloc = memref.alloc(%18) : memref<?xi8>
%c0_6 = arith.constant 0 : index
%c256_7 = arith.constant 256 : index
%c255_8 = arith.constant 255 : index
%19 = arith.addi %5, %c255_8 : index
%20 = arith.divui %19, %c256_7 : index
%21 = arith.muli %20, %c256_7 : index
%c1_9 = arith.constant 1 : index
%22 = arith.muli %21, %c1_9 : index
%23 = arith.addi %c0_6, %22 : index
%c2_10 = arith.constant 2 : index
%24 = arith.muli %21, %c2_10 : index
%25 = arith.addi %c0_6, %24 : index
%c3 = arith.constant 3 : index
%26 = arith.muli %21, %c3 : index
%27 = arith.addi %c0_6, %26 : index
%c4_11 = arith.constant 4 : index
%28 = arith.muli %21, %c4_11 : index
%29 = arith.addi %c0_6, %28 : index
%30 = arith.addi %8, %c255_8 : index
%31 = arith.divui %30, %c256_7 : index
%32 = arith.muli %31, %c256_7 : index
%c1_12 = arith.constant 1 : index
%33 = arith.muli %32, %c1_12 : index
%34 = arith.addi %29, %33 : index
%c2_13 = arith.constant 2 : index
%35 = arith.muli %32, %c2_13 : index
%36 = arith.addi %29, %35 : index
%view = memref.view %alloc[%c0_6][%dim, %dim_0] : memref<?xi8> to memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%view : memref<?x?x64xf32>)
%view_14 = memref.view %alloc[%23][%dim, %dim_0] : memref<?xi8> to memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %2 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%view_14 : memref<?x?x64xf32>)
%view_15 = memref.view %alloc[%25][%dim, %dim_0] : memref<?xi8> to memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %3 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%view_15 : memref<?x?x64xf32>)
%view_16 = memref.view %alloc[%27][%dim, %dim_0] : memref<?xi8> to memref<?x64x?xf32>
hip.transpose(%arg0, %c1, %c2) ins(%view_14 : memref<?x?x64xf32>) outs(%view_16 : memref<?x64x?xf32>)
%view_17 = memref.view %alloc[%29][%dim, %dim_0, %dim_0] : memref<?xi8> to memref<?x?x?xf32>
hip.hipblaslt.matmul(%arg0) ins(%view, %view_16 : memref<?x?x64xf32>, memref<?x64x?xf32>) outs(%view_17 : memref<?x?x?xf32>)
%view_18 = memref.view %alloc[%34][%dim, %dim_0, %dim_0] : memref<?xi8> to memref<?x?x?xf32>
hip.miopen.mul(%arg0) ins(%view_17, %0 : memref<?x?x?xf32>, memref<f32>) outs(%view_18 : memref<?x?x?xf32>)
hip.miopen.softmax(%arg0) ins(%view_18 : memref<?x?x?xf32>) outs(%view_17 : memref<?x?x?xf32>)
hip.hipblaslt.matmul(%arg0) ins(%view_17, %view_15 : memref<?x?x?xf32>, memref<?x?x64xf32>) outs(%arg2 : memref<?x?x64xf32>)
memref.dealloc %alloc : memref<?xi8>
return
}
Observation: Single memref.alloc / single memref.dealloc. All 6 buffers are memref.view slices into the pool. Pool size = 4 * align256(batch*seq*64*sizeof(f32)) + 2 * align256(batch*seq*seq*sizeof(f32)).
Passes 17-19: Post-pool cleanup (convert-bufferization-to-memref, cse, canonicalize)
Cleans up remaining bufferization ops, folds duplicate constants, and simplifies the pool offset arithmetic. After canonicalize, the pool size and offset computations are cleaner:
module attributes {hip.constants_file = "model.constants.bin", llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128", llvm.target_triple = "x86_64-unknown-linux-gnu", "onnx-mlir.symbol-postfix" = "attention_dynamic"} {
memref.global "private" constant @__constant_xf32 : memref<f32> = dense<1.250000e-01> {alignment = 64 : i64}
memref.global "private" @hip_ext_constant_2 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 32768 : i64, size = 16384 : i64}}
memref.global "private" @hip_ext_constant_1 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 16384 : i64, size = 16384 : i64}}
memref.global "private" @hip_ext_constant_0 : memref<64x64xf32> {alignment = 64 : i64, hip.external_data = {offset = 0 : i64, size = 16384 : i64}}
func.func @main_graph(%arg0: !hip.context, %arg1: memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>> {onnx.dim_params = "0:batch,1:seq", onnx.name = "X"}, %arg2: memref<?x?x64xf32> {bufferize.result, onnx.dim_params = "0:batch,1:seq", onnx.name = "out"}) attributes {hipdnn.buffer_offsets = [0, -1, -1, -1, -1, -1]} {
%c768 = arith.constant 768 : index
%c512 = arith.constant 512 : index
%c1024 = arith.constant 1024 : index
%c255 = arith.constant 255 : index
%c4 = arith.constant 4 : index
%c256 = arith.constant 256 : index
%c2 = arith.constant 2 : index
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%0 = memref.get_global @__constant_xf32 : memref<f32>
%1 = memref.get_global @hip_ext_constant_0 : memref<64x64xf32>
%2 = memref.get_global @hip_ext_constant_1 : memref<64x64xf32>
%3 = memref.get_global @hip_ext_constant_2 : memref<64x64xf32>
%dim = memref.dim %arg1, %c0 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%dim_0 = memref.dim %arg1, %c1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%4 = arith.muli %dim, %c256 : index
%5 = arith.muli %4, %dim_0 : index
%6 = arith.muli %dim, %c4 : index
%7 = arith.muli %6, %dim_0 : index
%8 = arith.muli %7, %dim_0 : index
%9 = arith.addi %5, %c255 : index
%10 = arith.divui %9, %c256 : index
%11 = arith.muli %10, %c256 : index
%12 = arith.muli %10, %c1024 : index
%13 = arith.addi %8, %c255 : index
%14 = arith.divui %13, %c256 : index
%15 = arith.muli %14, %c256 : index
%16 = arith.muli %14, %c512 : index
%17 = arith.addi %12, %16 : index
%alloc = memref.alloc(%17) : memref<?xi8>
%18 = arith.muli %10, %c512 : index
%19 = arith.muli %10, %c768 : index
%20 = arith.addi %12, %15 : index
%view = memref.view %alloc[%c0][%dim, %dim_0] : memref<?xi8> to memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%view : memref<?x?x64xf32>)
%view_1 = memref.view %alloc[%11][%dim, %dim_0] : memref<?xi8> to memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %2 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%view_1 : memref<?x?x64xf32>)
%view_2 = memref.view %alloc[%18][%dim, %dim_0] : memref<?xi8> to memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %3 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%view_2 : memref<?x?x64xf32>)
%view_3 = memref.view %alloc[%19][%dim, %dim_0] : memref<?xi8> to memref<?x64x?xf32>
hip.transpose(%arg0, %c1, %c2) ins(%view_1 : memref<?x?x64xf32>) outs(%view_3 : memref<?x64x?xf32>)
%view_4 = memref.view %alloc[%12][%dim, %dim_0, %dim_0] : memref<?xi8> to memref<?x?x?xf32>
hip.hipblaslt.matmul(%arg0) ins(%view, %view_3 : memref<?x?x64xf32>, memref<?x64x?xf32>) outs(%view_4 : memref<?x?x?xf32>)
%view_5 = memref.view %alloc[%20][%dim, %dim_0, %dim_0] : memref<?xi8> to memref<?x?x?xf32>
hip.miopen.mul(%arg0) ins(%view_4, %0 : memref<?x?x?xf32>, memref<f32>) outs(%view_5 : memref<?x?x?xf32>)
hip.miopen.softmax(%arg0) ins(%view_5 : memref<?x?x?xf32>) outs(%view_4 : memref<?x?x?xf32>)
hip.hipblaslt.matmul(%arg0) ins(%view_4, %view_2 : memref<?x?x?xf32>, memref<?x?x64xf32>) outs(%arg2 : memref<?x?x64xf32>)
memref.dealloc %alloc : memref<?xi8>
return
}
}
Pass 20: hip-lower-allocs
Replaces memref.alloc / memref.dealloc with hip.alloc / hip.free (GPU memory allocation via HIP runtime).
func.func @main_graph(%arg0: !hip.context, %arg1: memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>> {onnx.dim_params = "0:batch,1:seq", onnx.name = "X"}, %arg2: memref<?x?x64xf32> {bufferize.result, onnx.dim_params = "0:batch,1:seq", onnx.name = "out"}) attributes {hipdnn.buffer_offsets = [0, -1, -1, -1, -1, -1]} {
%c768 = arith.constant 768 : index
%c512 = arith.constant 512 : index
%c1024 = arith.constant 1024 : index
%c255 = arith.constant 255 : index
%c4 = arith.constant 4 : index
%c256 = arith.constant 256 : index
%c2 = arith.constant 2 : index
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%0 = memref.get_global @__constant_xf32 : memref<f32>
%1 = memref.get_global @hip_ext_constant_0 : memref<64x64xf32>
%2 = memref.get_global @hip_ext_constant_1 : memref<64x64xf32>
%3 = memref.get_global @hip_ext_constant_2 : memref<64x64xf32>
%dim = memref.dim %arg1, %c0 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%dim_0 = memref.dim %arg1, %c1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%4 = arith.muli %dim, %c256 : index
%5 = arith.muli %4, %dim_0 : index
%6 = arith.muli %dim, %c4 : index
%7 = arith.muli %6, %dim_0 : index
%8 = arith.muli %7, %dim_0 : index
%9 = arith.addi %5, %c255 : index
%10 = arith.divui %9, %c256 : index
%11 = arith.muli %10, %c256 : index
%12 = arith.muli %10, %c1024 : index
%13 = arith.addi %8, %c255 : index
%14 = arith.divui %13, %c256 : index
%15 = arith.muli %14, %c256 : index
%16 = arith.muli %14, %c512 : index
%17 = arith.addi %12, %16 : index
%18 = hip.alloc(%arg0, %17) : memref<?xi8>
%19 = arith.muli %10, %c512 : index
%20 = arith.muli %10, %c768 : index
%21 = arith.addi %12, %15 : index
%view = memref.view %18[%c0][%dim, %dim_0] : memref<?xi8> to memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%view : memref<?x?x64xf32>)
%view_1 = memref.view %18[%11][%dim, %dim_0] : memref<?xi8> to memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %2 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%view_1 : memref<?x?x64xf32>)
%view_2 = memref.view %18[%19][%dim, %dim_0] : memref<?xi8> to memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %3 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%view_2 : memref<?x?x64xf32>)
%view_3 = memref.view %18[%20][%dim, %dim_0] : memref<?xi8> to memref<?x64x?xf32>
hip.transpose(%arg0, %c1, %c2) ins(%view_1 : memref<?x?x64xf32>) outs(%view_3 : memref<?x64x?xf32>)
%view_4 = memref.view %18[%12][%dim, %dim_0, %dim_0] : memref<?xi8> to memref<?x?x?xf32>
hip.hipblaslt.matmul(%arg0) ins(%view, %view_3 : memref<?x?x64xf32>, memref<?x64x?xf32>) outs(%view_4 : memref<?x?x?xf32>)
%view_5 = memref.view %18[%21][%dim, %dim_0, %dim_0] : memref<?xi8> to memref<?x?x?xf32>
hip.miopen.mul(%arg0) ins(%view_4, %0 : memref<?x?x?xf32>, memref<f32>) outs(%view_5 : memref<?x?x?xf32>)
hip.miopen.softmax(%arg0) ins(%view_5 : memref<?x?x?xf32>) outs(%view_4 : memref<?x?x?xf32>)
hip.hipblaslt.matmul(%arg0) ins(%view_4, %view_2 : memref<?x?x?xf32>, memref<?x?x64xf32>) outs(%arg2 : memref<?x?x64xf32>)
hip.free(%arg0, %18) : memref<?xi8>
return
}
Observation: memref.alloc -> hip.alloc, memref.dealloc -> hip.free. Both take !hip.context as the first argument for GPU memory management.
Pass 21: hip-resolve-extern-constants
Resolves externalized constants. Removes memref.global ops with hip.external_data attributes and replaces memref.get_global with memref.view into a new constants buffer argument (%arg3: memref<?xi8>). The hip.constants_file module attribute is removed.
module attributes {llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128", llvm.target_triple = "x86_64-unknown-linux-gnu", "onnx-mlir.symbol-postfix" = "attention_dynamic"} {
memref.global "private" constant @__constant_xf32 : memref<f32> = dense<1.250000e-01> {alignment = 64 : i64}
func.func @main_graph(%arg0: !hip.context, %arg1: memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, %arg2: memref<?x?x64xf32>, %arg3: memref<?xi8>) attributes {hipdnn.buffer_offsets = [0, -1, -1, -1, -1, -1]} {
%c768 = arith.constant 768 : index
%c512 = arith.constant 512 : index
%c1024 = arith.constant 1024 : index
%c255 = arith.constant 255 : index
%c4 = arith.constant 4 : index
%c256 = arith.constant 256 : index
%c2 = arith.constant 2 : index
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%0 = memref.get_global @__constant_xf32 : memref<f32>
%c0_0 = arith.constant 0 : index
%view = memref.view %arg3[%c0_0][] : memref<?xi8> to memref<64x64xf32>
%c16384 = arith.constant 16384 : index
%view_1 = memref.view %arg3[%c16384][] : memref<?xi8> to memref<64x64xf32>
%c32768 = arith.constant 32768 : index
%view_2 = memref.view %arg3[%c32768][] : memref<?xi8> to memref<64x64xf32>
%dim = memref.dim %arg1, %c0 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%dim_3 = memref.dim %arg1, %c1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>
%1 = arith.muli %dim, %c256 : index
%2 = arith.muli %1, %dim_3 : index
%3 = arith.muli %dim, %c4 : index
%4 = arith.muli %3, %dim_3 : index
%5 = arith.muli %4, %dim_3 : index
%6 = arith.addi %2, %c255 : index
%7 = arith.divui %6, %c256 : index
%8 = arith.muli %7, %c256 : index
%9 = arith.muli %7, %c1024 : index
%10 = arith.addi %5, %c255 : index
%11 = arith.divui %10, %c256 : index
%12 = arith.muli %11, %c256 : index
%13 = arith.muli %11, %c512 : index
%14 = arith.addi %9, %13 : index
%15 = hip.alloc(%arg0, %14) : memref<?xi8>
%16 = arith.muli %7, %c512 : index
%17 = arith.muli %7, %c768 : index
%18 = arith.addi %9, %12 : index
%view_4 = memref.view %15[%c0][%dim, %dim_3] : memref<?xi8> to memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %view : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%view_4 : memref<?x?x64xf32>)
%view_5 = memref.view %15[%8][%dim, %dim_3] : memref<?xi8> to memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %view_1 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%view_5 : memref<?x?x64xf32>)
%view_6 = memref.view %15[%16][%dim, %dim_3] : memref<?xi8> to memref<?x?x64xf32>
hip.hipblaslt.matmul(%arg0) ins(%arg1, %view_2 : memref<?x?x64xf32, strided<[?, ?, ?], offset: ?>>, memref<64x64xf32>) outs(%view_6 : memref<?x?x64xf32>)
%view_7 = memref.view %15[%17][%dim, %dim_3] : memref<?xi8> to memref<?x64x?xf32>
hip.transpose(%arg0, %c1, %c2) ins(%view_5 : memref<?x?x64xf32>) outs(%view_7 : memref<?x64x?xf32>)
%view_8 = memref.view %15[%9][%dim, %dim_3, %dim_3] : memref<?xi8> to memref<?x?x?xf32>
hip.hipblaslt.matmul(%arg0) ins(%view_4, %view_7 : memref<?x?x64xf32>, memref<?x64x?xf32>) outs(%view_8 : memref<?x?x?xf32>)
%view_9 = memref.view %15[%18][%dim, %dim_3, %dim_3] : memref<?xi8> to memref<?x?x?xf32>
hip.miopen.mul(%arg0) ins(%view_8, %0 : memref<?x?x?xf32>, memref<f32>) outs(%view_9 : memref<?x?x?xf32>)
hip.miopen.softmax(%arg0) ins(%view_9 : memref<?x?x?xf32>) outs(%view_8 : memref<?x?x?xf32>)
hip.hipblaslt.matmul(%arg0) ins(%view_8, %view_6 : memref<?x?x?xf32>, memref<?x?x64xf32>) outs(%arg2 : memref<?x?x64xf32>)
hip.free(%arg0, %15) : memref<?xi8>
return
}
}
Observation: Final function signature: @main_graph(%ctx, %input, %output, %constants_blob). Constants are views into %arg3 at known offsets (0, 16384, 32768). Intermediates use a single hip.alloc/hip.free pool. Result: 7 dynamic GPU allocations reduced to 1.
Model
Self-attention with dynamic batch and sequence dimensions (
tensor<?x?x64xf32>), 3 weight projections (Q/K/V), scaled dot-product attention with softmax.Pipeline
hip-mlir-opt attention.onnx.mlir \ --allow-unregistered-dialect \ --onnx-to-hip-pipeline='externalize-min-num-elements=256 externalize-output-dir=/tmp' \ --mlir-elide-elementsattrs-if-larger=4 \ --mlir-print-ir-after-all21 passes executed in the
onnx-to-hip-pipeline. Complete IR after each pass is shown below.Pass 1:
hip-add-context-argInserts
!hip.contextas the first argument to@main_graph. The ONNX ops are unchanged.Pass 2:
convert-onnx-to-hipConverts ONNX ops to HIP dialect ops (tensor mode). Large constants (>= 256 elements) are externalized to
/tmp/model.constants.binand replaced withmemref.global+hip.external_dataattributes. Small constants are kept inline asarith.constant.Pass 3:
one-shot-bufferizeConverts tensor operations to memref operations.
tensor.emptybecomesmemref.alloc,tensor.dimbecomesmemref.dim. Function signature changes from tensor returns to memref returns. 8memref.alloccalls are introduced.Pass 4:
buffer-results-to-out-paramsConverts the memref return value to an output parameter (
%arg2). The function now returns void. The final matmul writes directly to%arg2instead of a locally-allocated buffer.Pass 5:
expand-reallocNo-op for this model (no realloc operations present). IR is identical to Pass 4.
Pass 6:
canonicalizeFolds redundant
memref.dimoperations (dim_1/dim_2 CSE'd with dim/dim_0), removes unusedmemref.cast, simplifies dim queries. Reduces from 8 to 7memref.alloc(the output alloc was already replaced by%arg2).Passes 7-14: Buffer deallocation and cleanup
Passes 7-14 are the standard MLIR buffer deallocation pipeline:
ownership-based-buffer-deallocation->canonicalize->buffer-deallocation-simplification->bufferization-lower-deallocations->cse->canonicalize->cse->canonicalize. The final result after pass 14 is the clean form with 7memref.alloc/ 7memref.deallocpairs:Pass 15:
hip-optimize-memrefs(liveness analysis)Analyzes buffer liveness and reuses dead buffers.
alloc_6(softmax output) is eliminated -- softmax now writes directly toalloc_4(which is dead after the QK^T matmul result is consumed by mul). 7 allocs reduced to 6.Pass 16:
hip-pool-allocs(memory pooling)All 6 dynamic allocations consolidated into a single
memref<?xi8>pool. Each buffer becomes amemref.viewat a 256-byte aligned offset. Pool size is computed dynamically at runtime.Passes 17-19: Post-pool cleanup (
convert-bufferization-to-memref,cse,canonicalize)Cleans up remaining bufferization ops, folds duplicate constants, and simplifies the pool offset arithmetic. After canonicalize, the pool size and offset computations are cleaner:
Pass 20:
hip-lower-allocsReplaces
memref.alloc/memref.deallocwithhip.alloc/hip.free(GPU memory allocation via HIP runtime).Pass 21:
hip-resolve-extern-constantsResolves externalized constants. Removes
memref.globalops withhip.external_dataattributes and replacesmemref.get_globalwithmemref.viewinto a new constants buffer argument (%arg3: memref<?xi8>). Thehip.constants_filemodule attribute is removed.