Skip to content

Instantly share code, notes, and snippets.

@pashu123

pashu123/all.txt

Created June 18, 2024 16:41
Show Gist options
  • Save pashu123/bb99ebc351d6cd9014c9315a361ecc4e to your computer and use it in GitHub Desktop.
Save pashu123/bb99ebc351d6cd9014c9315a361ecc4e to your computer and use it in GitHub Desktop.
Download ZIP
Raw
all.txt
This file has been truncated, but you can view the full file.
// -----// IR Dump After AssignTargetDevicesPass (iree-hal-assign-target-devices) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: tensor<?x?x3200xf32>, %arg1: tensor<8640x3200xf16>) -> tensor<?x?x8640xf32> {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%dim = tensor.dim %arg0, %c0 : tensor<?x?x3200xf32>
%dim_0 = tensor.dim %arg0, %c1 : tensor<?x?x3200xf32>
%0 = tensor.empty(%dim) : tensor<?x8640x3200xf16>
%1 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<8640x3200xf16>) outs(%0 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%2 = tensor.empty(%dim, %dim_0) : tensor<?x?x8640xf32>
%3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%4 = linalg.batch_matmul_transpose_b ins(%arg0, %1 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%3 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
util.return %4 : tensor<?x?x8640xf32>
}
}
// -----// IR Dump After AutoInputConversionPipeline (iree-auto-input-conversion) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: tensor<?x?x3200xf32>, %arg1: tensor<8640x3200xf16>) -> tensor<?x?x8640xf32> {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%dim = tensor.dim %arg0, %c0 : tensor<?x?x3200xf32>
%dim_0 = tensor.dim %arg0, %c1 : tensor<?x?x3200xf32>
%0 = tensor.empty(%dim) : tensor<?x8640x3200xf16>
%1 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<8640x3200xf16>) outs(%0 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%2 = tensor.empty(%dim, %dim_0) : tensor<?x?x8640xf32>
%3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%4 = linalg.batch_matmul_transpose_b ins(%arg0, %1 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%3 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
util.return %4 : tensor<?x?x8640xf32>
}
}
// -----// IR Dump After IREEImportPublic (iree-import-public) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: tensor<?x?x3200xf32>, %arg1: tensor<8640x3200xf16>) -> tensor<?x?x8640xf32> {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%dim = tensor.dim %arg0, %c0 : tensor<?x?x3200xf32>
%dim_0 = tensor.dim %arg0, %c1 : tensor<?x?x3200xf32>
%0 = tensor.empty(%dim) : tensor<?x8640x3200xf16>
%1 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<8640x3200xf16>) outs(%0 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%2 = tensor.empty(%dim, %dim_0) : tensor<?x?x8640xf32>
%3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%4 = linalg.batch_matmul_transpose_b ins(%arg0, %1 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%3 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
util.return %4 : tensor<?x?x8640xf32>
}
}
// -----// IR Dump After ImportMLProgram (iree-import-ml-program) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: tensor<?x?x3200xf32>, %arg1: tensor<8640x3200xf16>) -> tensor<?x?x8640xf32> {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%dim = tensor.dim %arg0, %c0 : tensor<?x?x3200xf32>
%dim_0 = tensor.dim %arg0, %c1 : tensor<?x?x3200xf32>
%0 = tensor.empty(%dim) : tensor<?x8640x3200xf16>
%1 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<8640x3200xf16>) outs(%0 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%2 = tensor.empty(%dim, %dim_0) : tensor<?x?x8640xf32>
%3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%4 = linalg.batch_matmul_transpose_b ins(%arg0, %1 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%3 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
util.return %4 : tensor<?x?x8640xf32>
}
}
// -----// IR Dump After SanitizeModuleNames (iree-sanitize-module-names) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: tensor<?x?x3200xf32>, %arg1: tensor<8640x3200xf16>) -> tensor<?x?x8640xf32> {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%dim = tensor.dim %arg0, %c0 : tensor<?x?x3200xf32>
%dim_0 = tensor.dim %arg0, %c1 : tensor<?x?x3200xf32>
%0 = tensor.empty(%dim) : tensor<?x8640x3200xf16>
%1 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<8640x3200xf16>) outs(%0 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%2 = tensor.empty(%dim, %dim_0) : tensor<?x?x8640xf32>
%3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%4 = linalg.batch_matmul_transpose_b ins(%arg0, %1 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%3 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
util.return %4 : tensor<?x?x8640xf32>
}
}
// -----// IR Dump After ConvertMeshToFlowPass (iree-convert-mesh-to-flow) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: tensor<?x?x3200xf32>, %arg1: tensor<8640x3200xf16>) -> tensor<?x?x8640xf32> {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%dim = tensor.dim %arg0, %c0 : tensor<?x?x3200xf32>
%dim_0 = tensor.dim %arg0, %c1 : tensor<?x?x3200xf32>
%0 = tensor.empty(%dim) : tensor<?x8640x3200xf16>
%1 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<8640x3200xf16>) outs(%0 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%2 = tensor.empty(%dim, %dim_0) : tensor<?x?x8640xf32>
%3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%4 = linalg.batch_matmul_transpose_b ins(%arg0, %1 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%3 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
util.return %4 : tensor<?x?x8640xf32>
}
}
// -----// IR Dump After DemoteF64ToF32 (iree-input-conversion-demote-f64-to-f32) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: tensor<?x?x3200xf32>, %arg1: tensor<8640x3200xf16>) -> tensor<?x?x8640xf32> {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%dim = tensor.dim %arg0, %c0 : tensor<?x?x3200xf32>
%dim_0 = tensor.dim %arg0, %c1 : tensor<?x?x3200xf32>
%0 = tensor.empty(%dim) : tensor<?x8640x3200xf16>
%1 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<8640x3200xf16>) outs(%0 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%2 = tensor.empty(%dim, %dim_0) : tensor<?x?x8640xf32>
%3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%4 = linalg.batch_matmul_transpose_b ins(%arg0, %1 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%3 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
util.return %4 : tensor<?x?x8640xf32>
}
}
// -----// IR Dump After mlir::iree_compiler::IREE::ABI::ConvertStreamableOpsPass (iree-abi-convert-streamable-ops) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: tensor<?x?x3200xf32>, %arg1: tensor<8640x3200xf16>) -> tensor<?x?x8640xf32> {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%dim = tensor.dim %arg0, %c0 : tensor<?x?x3200xf32>
%dim_0 = tensor.dim %arg0, %c1 : tensor<?x?x3200xf32>
%0 = tensor.empty(%dim) : tensor<?x8640x3200xf16>
%1 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<8640x3200xf16>) outs(%0 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%2 = tensor.empty(%dim, %dim_0) : tensor<?x?x8640xf32>
%3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%4 = linalg.batch_matmul_transpose_b ins(%arg0, %1 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%3 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
util.return %4 : tensor<?x?x8640xf32>
}
}
// -----// IR Dump After mlir::iree_compiler::IREE::ABI::WrapEntryPointsPass (iree-abi-wrap-entry-points) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = util.call @_matmul_broad(%2, %3) : (tensor<?x?x3200xf32>, tensor<8640x3200xf16>) -> tensor<?x?x8640xf32>
%c0 = arith.constant 0 : index
%dim = tensor.dim %4, %c0 : tensor<?x?x8640xf32>
%c1 = arith.constant 1 : index
%dim_0 = tensor.dim %4, %c1 : tensor<?x?x8640xf32>
%5 = hal.tensor.export %4 "output0" : tensor<?x?x8640xf32>{%dim, %dim_0} -> !hal.buffer_view
util.return %5 : !hal.buffer_view
}
util.func private @_matmul_broad(%arg0: tensor<?x?x3200xf32>, %arg1: tensor<8640x3200xf16>) -> tensor<?x?x8640xf32> {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%dim = tensor.dim %arg0, %c0 : tensor<?x?x3200xf32>
%dim_0 = tensor.dim %arg0, %c1 : tensor<?x?x3200xf32>
%0 = tensor.empty(%dim) : tensor<?x8640x3200xf16>
%1 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<8640x3200xf16>) outs(%0 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%2 = tensor.empty(%dim, %dim_0) : tensor<?x?x8640xf32>
%3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%4 = linalg.batch_matmul_transpose_b ins(%arg0, %1 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%3 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
util.return %4 : tensor<?x?x8640xf32>
}
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func private @_matmul_broad(%arg0: tensor<?x?x3200xf32>, %arg1: tensor<8640x3200xf16>) -> tensor<?x?x8640xf32> {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%dim = tensor.dim %arg0, %c0 : tensor<?x?x3200xf32>
%dim_0 = tensor.dim %arg0, %c1 : tensor<?x?x3200xf32>
%0 = tensor.empty(%dim) : tensor<?x8640x3200xf16>
%1 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<8640x3200xf16>) outs(%0 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%2 = tensor.empty(%dim, %dim_0) : tensor<?x?x8640xf32>
%3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%4 = linalg.batch_matmul_transpose_b ins(%arg0, %1 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%3 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
util.return %4 : tensor<?x?x8640xf32>
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = util.call @_matmul_broad(%2, %3) : (tensor<?x?x3200xf32>, tensor<8640x3200xf16>) -> tensor<?x?x8640xf32>
%dim = tensor.dim %4, %c0 : tensor<?x?x8640xf32>
%dim_0 = tensor.dim %4, %c1 : tensor<?x?x8640xf32>
%5 = hal.tensor.export %4 "output0" : tensor<?x?x8640xf32>{%dim, %dim_0} -> !hal.buffer_view
util.return %5 : !hal.buffer_view
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After Inliner (inline) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After SymbolDCE (symbol-dce) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After RemoveZeroExtentTensors (iree-global-opt-remove-zero-extent-tensors) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After DetachElementwiseFromNamedOps (iree-global-opt-detach-elementwise-from-named-ops) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After LinalgNamedOpConversionPass (linalg-named-op-conversion) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After Convert1X1FilterConv2DToMatmul (iree-global-opt-convert-1x1-filter-conv2d-to-matmul) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After EraseUnusedLinalgOperands (iree-global-opt-erase-unused-linalg-operands) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After ExpandTensorShapes (iree-global-opt-expand-tensor-shapes) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After ConvertElementwiseToLinalgPass (convert-elementwise-to-linalg) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After RaiseSpecialOps (iree-global-opt-raise-special-ops) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After DecomposeConcat (iree-global-opt-decompose-concat) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After GeneralizeLinalgNamedOps (iree-global-opt-generalize-linalg-named-ops) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After FoldUnitExtentDimsPass (iree-flow-fold-unit-extent-dims) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After DemoteContractionInputsToBF16 (iree-global-opt-demote-contraction-inputs-to-bf16) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After FuseDequantizationMatmul (iree-global-opt-fuse-dequantization-matmul) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%8 = linalg.batch_matmul_transpose_b ins(%2, %5 : tensor<?x?x3200xf32>, tensor<?x8640x3200xf16>) outs(%7 : tensor<?x?x8640xf32>) -> tensor<?x?x8640xf32>
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After SetEncoding (iree-global-opt-set-encoding) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
#map3 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
#map4 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
#map5 = affine_map<()[s0, s1] -> (-s1 + (s1 ceildiv s0) * s0)>
#map6 = affine_map<()[s0, s1, s2] -> (-s1 + s2 + (s1 ceildiv s0) * s0)>
#map7 = affine_map<()[s0] -> ((8640 ceildiv s0) * s0)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f16
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%cst_0 = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6:3 = iree_encoding.upper_bound_tile_size tensor<?x?x3200xf32, #iree_encoding.encoding<role = LHS, element_types = [f32, f16, f32], user_indexing_maps = [#map2, #map3, #map4]>> -> index, index, index
%dim = tensor.dim %2, %c0 : tensor<?x?x3200xf32>
%7 = affine.apply #map5()[%6#0, %dim]
%dim_1 = tensor.dim %2, %c1 : tensor<?x?x3200xf32>
%8 = affine.apply #map5()[%6#1, %dim_1]
%9 = affine.apply #map5()[%6#2, %c3200]
%padded = tensor.pad %2 low[0, 0, 0] high[%7, %8, %9] {
^bb0(%arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_0 : f32
} : tensor<?x?x3200xf32> to tensor<?x?x?xf32>
%10 = iree_encoding.set_encoding %padded : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #iree_encoding.encoding<role = LHS, element_types = [f32, f16, f32], original_type = tensor<?x?x3200xf32>, user_indexing_maps = [#map2, #map3, #map4]>>
%11:3 = iree_encoding.upper_bound_tile_size tensor<?x8640x3200xf16, #iree_encoding.encoding<role = RHS, element_types = [f32, f16, f32], user_indexing_maps = [#map2, #map3, #map4]>> -> index, index, index
%12 = affine.apply #map5()[%11#0, %0]
%13 = affine.apply #map5()[%11#1, %c8640]
%14 = affine.apply #map5()[%11#2, %c3200]
%padded_2 = tensor.pad %5 low[0, 0, 0] high[%12, %13, %14] {
^bb0(%arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst : f16
} : tensor<?x8640x3200xf16> to tensor<?x?x?xf16>
%15 = iree_encoding.set_encoding %padded_2 : tensor<?x?x?xf16> -> tensor<?x?x?xf16, #iree_encoding.encoding<role = RHS, element_types = [f32, f16, f32], original_type = tensor<?x8640x3200xf16>, user_indexing_maps = [#map2, #map3, #map4]>>
%16:3 = iree_encoding.upper_bound_tile_size tensor<?x?x8640xf32, #iree_encoding.encoding<role = RESULT, element_types = [f32, f16, f32], user_indexing_maps = [#map2, #map3, #map4]>> -> index, index, index
%17 = affine.apply #map6()[%16#0, %0, %0]
%18 = affine.apply #map6()[%16#1, %1, %1]
%19 = affine.apply #map7()[%16#2]
%20 = tensor.empty(%17, %18, %19) : tensor<?x?x?xf32, #iree_encoding.encoding<role = RESULT, element_types = [f32, f16, f32], original_type = tensor<?x?x8640xf32>, user_indexing_maps = [#map2, #map3, #map4]>>
%21 = linalg.fill ins(%cst_0 : f32) outs(%20 : tensor<?x?x?xf32, #iree_encoding.encoding<role = RESULT, element_types = [f32, f16, f32], original_type = tensor<?x?x8640xf32>, user_indexing_maps = [#map2, #map3, #map4]>>) -> tensor<?x?x?xf32, #iree_encoding.encoding<role = RESULT, element_types = [f32, f16, f32], original_type = tensor<?x?x8640xf32>, user_indexing_maps = [#map2, #map3, #map4]>>
%22 = linalg.batch_matmul_transpose_b ins(%10, %15 : tensor<?x?x?xf32, #iree_encoding.encoding<role = LHS, element_types = [f32, f16, f32], original_type = tensor<?x?x3200xf32>, user_indexing_maps = [#map2, #map3, #map4]>>, tensor<?x?x?xf16, #iree_encoding.encoding<role = RHS, element_types = [f32, f16, f32], original_type = tensor<?x8640x3200xf16>, user_indexing_maps = [#map2, #map3, #map4]>>) outs(%21 : tensor<?x?x?xf32, #iree_encoding.encoding<role = RESULT, element_types = [f32, f16, f32], original_type = tensor<?x?x8640xf32>, user_indexing_maps = [#map2, #map3, #map4]>>) -> tensor<?x?x?xf32, #iree_encoding.encoding<role = RESULT, element_types = [f32, f16, f32], original_type = tensor<?x?x8640xf32>, user_indexing_maps = [#map2, #map3, #map4]>>
%23 = iree_encoding.unset_encoding %22 : tensor<?x?x?xf32, #iree_encoding.encoding<role = RESULT, element_types = [f32, f16, f32], original_type = tensor<?x?x8640xf32>, user_indexing_maps = [#map2, #map3, #map4]>> -> tensor<?x?x?xf32>
%extracted_slice = tensor.extract_slice %23[0, 0, 0] [%0, %1, 8640] [1, 1, 1] : tensor<?x?x?xf32> to tensor<?x?x8640xf32>
%24 = hal.tensor.export %extracted_slice "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %24 : !hal.buffer_view
}
}
innerTileSizes: 8
innerTileSizes: 4
innerDimsPos: 1
innerDimsPos: 2
innerTileSizes: 4
innerTileSizes: 1
innerDimsPos: 1
innerDimsPos: 2
innerTileSizes: 8
innerTileSizes: 1
innerDimsPos: 1
innerDimsPos: 2
// -----// IR Dump After CPUMaterializeUpperBoundTileSize (iree-codegen-cpu-materialize-upper-bound-tile-size) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c8640 = arith.constant 8640 : index
%c8 = arith.constant 8 : index
%cst = arith.constant 0.000000e+00 : f16
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%cst_0 = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%dim = tensor.dim %2, %c0 : tensor<?x?x3200xf32>
%6 = affine.apply affine_map<()[s0, s1] -> (-s1 + (s1 ceildiv s0) * s0)>()[%c1, %dim]
%dim_1 = tensor.dim %2, %c1 : tensor<?x?x3200xf32>
%7 = affine.apply affine_map<()[s0, s1] -> (-s1 + (s1 ceildiv s0) * s0)>()[%c8, %dim_1]
%padded = tensor.pad %2 low[0, 0, 0] high[%6, %7, %c0] {
^bb0(%arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst_0 : f32
} : tensor<?x?x3200xf32> to tensor<?x?x?xf32>
%8 = iree_encoding.set_encoding %padded : tensor<?x?x?xf32> -> tensor<?x?x?xf32, #iree_encoding.encoding<role = LHS, element_types = [f32, f16, f32], original_type = tensor<?x?x3200xf32>, user_indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>]>>
%9 = affine.apply affine_map<()[s0, s1] -> (-s1 + (s1 ceildiv s0) * s0)>()[%c1, %0]
%padded_2 = tensor.pad %5 low[0, 0, 0] high[%9, %c0, %c0] {
^bb0(%arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst : f16
} : tensor<?x8640x3200xf16> to tensor<?x?x?xf16>
%10 = iree_encoding.set_encoding %padded_2 : tensor<?x?x?xf16> -> tensor<?x?x?xf16, #iree_encoding.encoding<role = RHS, element_types = [f32, f16, f32], original_type = tensor<?x8640x3200xf16>, user_indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>]>>
%11 = affine.apply affine_map<()[s0, s1, s2] -> (-s1 + s2 + (s1 ceildiv s0) * s0)>()[%c1, %0, %0]
%12 = affine.apply affine_map<()[s0, s1, s2] -> (-s1 + s2 + (s1 ceildiv s0) * s0)>()[%c8, %1, %1]
%13 = tensor.empty(%11, %12, %c8640) : tensor<?x?x?xf32, #iree_encoding.encoding<role = RESULT, element_types = [f32, f16, f32], original_type = tensor<?x?x8640xf32>, user_indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>]>>
%14 = linalg.fill ins(%cst_0 : f32) outs(%13 : tensor<?x?x?xf32, #iree_encoding.encoding<role = RESULT, element_types = [f32, f16, f32], original_type = tensor<?x?x8640xf32>, user_indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>]>>) -> tensor<?x?x?xf32, #iree_encoding.encoding<role = RESULT, element_types = [f32, f16, f32], original_type = tensor<?x?x8640xf32>, user_indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>]>>
%15 = linalg.batch_matmul_transpose_b ins(%8, %10 : tensor<?x?x?xf32, #iree_encoding.encoding<role = LHS, element_types = [f32, f16, f32], original_type = tensor<?x?x3200xf32>, user_indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>]>>, tensor<?x?x?xf16, #iree_encoding.encoding<role = RHS, element_types = [f32, f16, f32], original_type = tensor<?x8640x3200xf16>, user_indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>]>>) outs(%14 : tensor<?x?x?xf32, #iree_encoding.encoding<role = RESULT, element_types = [f32, f16, f32], original_type = tensor<?x?x8640xf32>, user_indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>]>>) -> tensor<?x?x?xf32, #iree_encoding.encoding<role = RESULT, element_types = [f32, f16, f32], original_type = tensor<?x?x8640xf32>, user_indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>]>>
%16 = iree_encoding.unset_encoding %15 : tensor<?x?x?xf32, #iree_encoding.encoding<role = RESULT, element_types = [f32, f16, f32], original_type = tensor<?x?x8640xf32>, user_indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>]>> -> tensor<?x?x?xf32>
%extracted_slice = tensor.extract_slice %16[0, 0, 0] [%0, %1, 8640] [1, 1, 1] : tensor<?x?x?xf32> to tensor<?x?x8640xf32>
%17 = hal.tensor.export %extracted_slice "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %17 : !hal.buffer_view
}
%pack = tensor.pack %2 padding_value(%cst_0 : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %9 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%pack_9 = tensor.pack %5 padding_value(%cst : f16) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [4, 1] into %12 : tensor<?x8640x3200xf16> -> tensor<?x2160x3200x4x1xf16>
// -----// IR Dump After CPUMaterializeEncoding (iree-codegen-cpu-materialize-encoding) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f16
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%cst_0 = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%dim = tensor.dim %2, %c0 : tensor<?x?x3200xf32>
%dim_1 = tensor.dim %2, %c1 : tensor<?x?x3200xf32>
%6 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%dim_1]
%7 = tensor.empty(%dim, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %2 padding_value(%cst_0 : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%pack_2 = tensor.pack %5 padding_value(%cst : f16) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [4, 1] into %8 : tensor<?x8640x3200xf16> -> tensor<?x2160x3200x4x1xf16>
%9 = affine.apply affine_map<()[s0, s1, s2] -> (-s1 + s2 + (s1 ceildiv s0) * s0)>()[%c1, %0, %0]
%10 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%11 = tensor.empty(%9, %10) : tensor<?x?x2160x8x4xf32>
%12 = linalg.fill ins(%cst_0 : f32) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%13 = linalg.batch_mmt4d ins(%pack, %pack_2 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%12 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%14 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %13 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %14 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%15 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %15 : !hal.buffer_view
}
// -----// IR Dump After MaterializeHomogeneousEncodings (iree-global-opt-materialize-homogeneous-encodings) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#map3 = affine_map<()[s0, s1, s2] -> (-s1 + s2 + (s1 ceildiv s0) * s0)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f16
%c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index
%cst_0 = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%dim = tensor.dim %2, %c0 : tensor<?x?x3200xf32>
%dim_1 = tensor.dim %2, %c1 : tensor<?x?x3200xf32>
%6 = affine.apply #map2()[%dim_1]
%7 = tensor.empty(%dim, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %2 padding_value(%cst_0 : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%pack_2 = tensor.pack %5 padding_value(%cst : f16) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [4, 1] into %8 : tensor<?x8640x3200xf16> -> tensor<?x2160x3200x4x1xf16>
%9 = affine.apply #map3()[%c1, %0, %0]
%10 = affine.apply #map2()[%1]
%11 = tensor.empty(%9, %10) : tensor<?x?x2160x8x4xf32>
%12 = linalg.fill ins(%cst_0 : f32) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%13 = linalg.batch_mmt4d ins(%pack, %pack_2 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%12 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%14 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %13 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %14 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%15 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %15 : !hal.buffer_view
}
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = affine.apply #map2()[%1]
%7 = tensor.empty(%0, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%pack_0 = tensor.pack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [4, 1] into %8 : tensor<?x8640x3200xf16> -> tensor<?x2160x3200x4x1xf16>
%9 = affine.apply #map2()[%1]
%10 = tensor.empty(%0, %9) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.batch_mmt4d ins(%pack, %pack_0 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%13 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %12 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %13 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%14 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %14 : !hal.buffer_view
}
}
// -----// IR Dump After CSE (cse) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = affine.apply #map2()[%1]
%7 = tensor.empty(%0, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%pack_0 = tensor.pack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [4, 1] into %8 : tensor<?x8640x3200xf16> -> tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %6) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack, %pack_0 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
}
// -----// IR Dump After SimplifyPackUnpack (iree-global-opt-simplify-pack-unpack) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x8640x3200xf16>
%5 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%6 = affine.apply #map2()[%1]
%7 = tensor.empty(%0, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%pack_0 = tensor.pack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [4, 1] into %8 : tensor<?x8640x3200xf16> -> tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %6) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack, %pack_0 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
}
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
Hey I am hereHey I am here%5 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%3 : tensor<8640x3200xf16>) outs(%4 : tensor<?x8640x3200xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x8640x3200xf16>
%8 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack_0 : tensor<2160x3200x4x1xf16>) outs(%5 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
Hey I am hereHey I am hereHey I am hereHey I am here%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
Hey I am hereHey I am here%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
Hey I am hereHey I am here%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
// -----// IR Dump After DataLayoutPropagation (iree-global-opt-data-layout-propagation) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After GeneralizeLinalgNamedOps (iree-global-opt-generalize-linalg-named-ops) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After GlobalLoopInvariantCodeMotion (iree-global-opt-loop-invariant-code-motion) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After SimplifyGlobalAccesses (iree-util-simplify-global-accesses) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After ApplyPatterns (iree-util-apply-patterns) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply #map2()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
}
// -----// IR Dump After FoldGlobals (iree-util-fold-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply #map2()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
}
// -----// IR Dump After IPO (iree-util-ipo) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply #map2()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply #map2()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
}
// -----// IR Dump After CSE (cse) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply #map2()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
}
// -----// IR Dump After HoistIntoGlobals (iree-util-hoist-into-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply #map2()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
}
// -----// IR Dump After JitGlobals (iree-consteval-jit-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply #map2()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After RaiseSpecialOps (iree-global-opt-raise-special-ops) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After VerifyInputLegalityPass (iree-verify-input-legality) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply #map2()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
}
// -----// IR Dump After InjectTensorTracingPass (iree-flow-inject-tensor-tracing) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After TensorPadToTensorInsertSlicePass (iree-flow-tensor-pad-to-tensor-insert-slice) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply #map2()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After SimplifyGlobalAccesses (iree-util-simplify-global-accesses) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After ApplyPatterns (iree-util-apply-patterns) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local], iree.fixedpoint.iteration = 0 : index} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply #map2()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
}
// -----// IR Dump After FoldGlobals (iree-util-fold-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local], iree.fixedpoint.iteration = 0 : index} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply #map2()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
}
// -----// IR Dump After FuseGlobals (iree-util-fuse-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local], iree.fixedpoint.iteration = 0 : index} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply #map2()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
}
// -----// IR Dump After IPO (iree-util-ipo) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local], iree.fixedpoint.iteration = 0 : index} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply #map2()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
}
// -----// IR Dump After FixedPointIterator (iree-util-fixed-point-iterator) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#map2 = affine_map<()[s0] -> (s0 ceildiv 8)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply #map2()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
}
// -----// IR Dump After FusionPreprocessingPass (iree-flow-fusion-preprocessing) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After ElementwiseOpFusionPass (iree-flow-elementwise-op-fusion) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After BubbleUpExpandShapesPass (iree-flow-bubble-up-expand-shapes) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After ElementwiseOpFusionPass (iree-flow-elementwise-op-fusion) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After SinkReshapesPass (iree-flow-sink-reshapes) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After FusionOfTensorOpsPass (iree-flow-fusion-of-tensor-ops) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After CollapseReductionDimensionsPass (iree-flow-collapse-reduction-dimensions) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After SplitReductionPass (iree-flow-split-reduction-ops) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After FusionPreprocessingPass (iree-flow-fusion-preprocessing) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After InterchangeTransposeGenericOpsPass (iree-flow-interchange-transpose-generic-ops) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After FormScalarDispatchesPass (iree-flow-form-scalar-dispatches) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
%6 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%pack : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%pack_0 = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.batch_mmt4d ins(%pack_0, %6 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %11 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %12 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
%13 = hal.tensor.export %unpack "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %13 : !hal.buffer_view
}
// -----// IR Dump After FormDispatchRegionsPass (iree-flow-form-dispatch-regions) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.region -> (tensor<2160x3200x4x1xf16>) {
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %5 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.return %pack : tensor<2160x3200x4x1xf16>
}
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%9 = flow.dispatch.region -> (tensor<?x?x3200x8x1xf32>{%0, %7}) {
%pack = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %8 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.return %pack : tensor<?x?x3200x8x1xf32>
}
%10 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = flow.dispatch.region -> (tensor<?x?x2160x8x4xf32>{%0, %7}) {
%16 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%4 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%17 = linalg.batch_mmt4d ins(%9, %16 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.return %17 : tensor<?x?x2160x8x4xf32>
}
%13 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%c0_0 = arith.constant 0 : index
%c1_1 = arith.constant 1 : index
%14 = flow.dispatch.region -> (tensor<?x?x8640xf32>{%0, %1}) {
%unpack = tensor.unpack %12 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %13 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.return %unpack : tensor<?x?x8640xf32>
}
%15 = hal.tensor.export %14 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %15 : !hal.buffer_view
}
// -----// IR Dump After CloneProducersIntoDispatchRegionsPass (iree-flow-clone-producers-into-dispatch-regions) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.region -> (tensor<2160x3200x4x1xf16>) {
%16 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %16 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.return %pack : tensor<2160x3200x4x1xf16>
}
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%9 = flow.dispatch.region -> (tensor<?x?x3200x8x1xf32>{%0, %7}) {
%cst_2 = arith.constant 0.000000e+00 : f32
%16 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%17 = tensor.empty(%0, %16) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %2 padding_value(%cst_2 : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %17 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.return %pack : tensor<?x?x3200x8x1xf32>
}
%10 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = flow.dispatch.region -> (tensor<?x?x2160x8x4xf32>{%0, %7}) {
%16 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%cst_2 = arith.constant 0.000000e+00 : f32
%17 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%18 = tensor.empty(%0, %16) : tensor<?x?x2160x8x4xf32>
%19 = linalg.fill ins(%cst_2 : f32) outs(%18 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%20 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%17 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%21 = linalg.batch_mmt4d ins(%9, %20 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%19 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.return %21 : tensor<?x?x2160x8x4xf32>
}
%13 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%c0_0 = arith.constant 0 : index
%c1_1 = arith.constant 1 : index
%14 = flow.dispatch.region -> (tensor<?x?x8640xf32>{%0, %1}) {
%16 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %12 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %16 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.return %unpack : tensor<?x?x8640xf32>
}
%15 = hal.tensor.export %14 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %15 : !hal.buffer_view
}
// -----// IR Dump After CollapseDimensionsPass (iree-flow-collapse-dimensions) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%cst = arith.constant 0.000000e+00 : f32
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%5 = tensor.empty() : tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.region -> (tensor<2160x3200x4x1xf16>) {
%16 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %3 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %16 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.return %pack : tensor<2160x3200x4x1xf16>
}
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = tensor.empty(%0, %7) : tensor<?x?x3200x8x1xf32>
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%9 = flow.dispatch.region -> (tensor<?x?x3200x8x1xf32>{%0, %7}) {
%cst_2 = arith.constant 0.000000e+00 : f32
%16 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%17 = tensor.empty(%0, %16) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %2 padding_value(%cst_2 : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %17 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.return %pack : tensor<?x?x3200x8x1xf32>
}
%10 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = flow.dispatch.region -> (tensor<?x?x2160x8x4xf32>{%0, %7}) {
%16 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%cst_2 = arith.constant 0.000000e+00 : f32
%17 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%18 = tensor.empty(%0, %16) : tensor<?x?x2160x8x4xf32>
%19 = linalg.fill ins(%cst_2 : f32) outs(%18 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%20 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%17 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%21 = linalg.batch_mmt4d ins(%9, %20 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%19 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.return %21 : tensor<?x?x2160x8x4xf32>
}
%13 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%c0_0 = arith.constant 0 : index
%c1_1 = arith.constant 1 : index
%14 = flow.dispatch.region -> (tensor<?x?x8640xf32>{%0, %1}) {
%16 = tensor.empty(%0, %1) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %12 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %16 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.return %unpack : tensor<?x?x8640xf32>
}
%15 = hal.tensor.export %14 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %15 : !hal.buffer_view
}
// -----// IR Dump After FormDispatchWorkgroupsPass (iree-flow-form-dispatch-workgroups) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch.workgroups(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16> =
(%arg2: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg3: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%10 = flow.dispatch.tensor.load %arg2, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%11 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %10 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %11 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg3, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
flow.return
} count() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
%5 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%6 = flow.dispatch.workgroups[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5} =
(%arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg3: index, %arg4: index, %arg5: index, %arg6: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%10 = flow.dispatch.workload.ordinal %arg3, 0 : index
%11 = flow.dispatch.workload.ordinal %arg4, 1 : index
%12 = flow.dispatch.workload.ordinal %arg5, 2 : index
%cst = arith.constant 0.000000e+00 : f32
%13 = flow.dispatch.tensor.load %arg2, offsets = [0, 0, 0], sizes = [%11, %10, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%11, %10} -> tensor<?x?x3200xf32>
%14 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%10]
%15 = tensor.empty(%11, %14) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %13 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %15 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %arg6, offsets = [0, 0, 0, 0, 0], sizes = [%11, %12, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%11, %12}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%7 = flow.dispatch.workgroups[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5} =
(%arg2: index, %arg3: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg4: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg5: index, %arg6: index, %arg7: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%10 = flow.dispatch.workload.ordinal %arg2, 0 : index
%11 = flow.dispatch.workload.ordinal %arg5, 1 : index
%12 = flow.dispatch.workload.ordinal %arg6, 2 : index
%cst = arith.constant 0.000000e+00 : f32
%13 = flow.dispatch.tensor.load %arg3, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%14 = flow.dispatch.tensor.load %arg4, offsets = [0, 0, 0, 0, 0], sizes = [%11, %12, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%11, %12} -> tensor<?x?x3200x8x1xf32>
%15 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%10]
%16 = tensor.empty(%11) : tensor<?x2160x3200x4x1xf16>
%17 = tensor.empty(%11, %15) : tensor<?x?x2160x8x4xf32>
%18 = linalg.fill ins(%cst : f32) outs(%17 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%19 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%13 : tensor<2160x3200x4x1xf16>) outs(%16 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%20 = linalg.batch_mmt4d ins(%14, %19 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%18 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %20, %arg7, offsets = [0, 0, 0, 0, 0], sizes = [%11, %12, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%11, %12}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%8 = flow.dispatch.workgroups[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1} =
(%arg2: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg3: index, %arg4: index, %arg5: index, %arg6: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%10 = flow.dispatch.workload.ordinal %arg3, 0 : index
%11 = flow.dispatch.workload.ordinal %arg4, 1 : index
%12 = flow.dispatch.workload.ordinal %arg5, 2 : index
%13 = flow.dispatch.tensor.load %arg2, offsets = [0, 0, 0, 0, 0], sizes = [%11, %10, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%11, %10} -> tensor<?x?x2160x8x4xf32>
%14 = tensor.empty(%11, %12) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %13 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %14 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %arg6, offsets = [0, 0, 0], sizes = [%11, %12, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%11, %12}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After CaptureDynamicDimsPass (iree-flow-capture-dynamic-dims) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch.workgroups(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16> =
(%arg2: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg3: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%10 = flow.dispatch.tensor.load %arg2, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%11 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %10 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %11 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg3, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
flow.return
} count() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
%5 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%6 = flow.dispatch.workgroups[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5} =
(%arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg3: index, %arg4: index, %arg5: index, %arg6: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%10 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%arg4, %arg3}
%11 = flow.dispatch.tie_shape %arg6 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%arg4, %arg5}
%12 = flow.dispatch.workload.ordinal %arg3, 0 : index
%13 = flow.dispatch.workload.ordinal %arg4, 1 : index
%14 = flow.dispatch.workload.ordinal %arg5, 2 : index
%cst = arith.constant 0.000000e+00 : f32
%15 = flow.dispatch.tensor.load %10, offsets = [0, 0, 0], sizes = [%13, %12, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%13, %12} -> tensor<?x?x3200xf32>
%16 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%12]
%17 = tensor.empty(%13, %16) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %15 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %17 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %11, offsets = [0, 0, 0, 0, 0], sizes = [%13, %14, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%13, %14}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%7 = flow.dispatch.workgroups[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5} =
(%arg2: index, %arg3: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg4: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg5: index, %arg6: index, %arg7: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%10 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%arg5, %arg6}
%11 = flow.dispatch.tie_shape %arg7 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%arg5, %arg6}
%12 = flow.dispatch.workload.ordinal %arg2, 0 : index
%13 = flow.dispatch.workload.ordinal %arg5, 1 : index
%14 = flow.dispatch.workload.ordinal %arg6, 2 : index
%cst = arith.constant 0.000000e+00 : f32
%15 = flow.dispatch.tensor.load %arg3, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%16 = flow.dispatch.tensor.load %10, offsets = [0, 0, 0, 0, 0], sizes = [%13, %14, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%13, %14} -> tensor<?x?x3200x8x1xf32>
%17 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%12]
%18 = tensor.empty(%13) : tensor<?x2160x3200x4x1xf16>
%19 = tensor.empty(%13, %17) : tensor<?x?x2160x8x4xf32>
%20 = linalg.fill ins(%cst : f32) outs(%19 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%21 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%15 : tensor<2160x3200x4x1xf16>) outs(%18 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%22 = linalg.batch_mmt4d ins(%16, %21 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%20 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %22, %11, offsets = [0, 0, 0, 0, 0], sizes = [%13, %14, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%13, %14}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%8 = flow.dispatch.workgroups[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1} =
(%arg2: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg3: index, %arg4: index, %arg5: index, %arg6: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%10 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%arg4, %arg3}
%11 = flow.dispatch.tie_shape %arg6 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%arg4, %arg5}
%12 = flow.dispatch.workload.ordinal %arg3, 0 : index
%13 = flow.dispatch.workload.ordinal %arg4, 1 : index
%14 = flow.dispatch.workload.ordinal %arg5, 2 : index
%15 = flow.dispatch.tensor.load %10, offsets = [0, 0, 0, 0, 0], sizes = [%13, %12, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%13, %12} -> tensor<?x?x2160x8x4xf32>
%16 = tensor.empty(%13, %14) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %15 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %16 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %11, offsets = [0, 0, 0], sizes = [%13, %14, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%13, %14}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch.workgroups(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16> =
(%arg2: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg3: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%10 = flow.dispatch.tensor.load %arg2, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%11 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %10 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %11 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg3, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
flow.return
} count() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
%5 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%6 = flow.dispatch.workgroups[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5} =
(%arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg3: index, %arg4: index, %arg5: index, %arg6: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%10 = flow.dispatch.workload.ordinal %arg3, 0 : index
%11 = flow.dispatch.workload.ordinal %arg4, 1 : index
%12 = flow.dispatch.workload.ordinal %arg5, 2 : index
%13 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%11, %10}
%14 = flow.dispatch.tie_shape %arg6 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%11, %12}
%15 = flow.dispatch.tensor.load %13, offsets = [0, 0, 0], sizes = [%11, %10, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%11, %10} -> tensor<?x?x3200xf32>
%16 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%10]
%17 = tensor.empty(%11, %16) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %15 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %17 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %14, offsets = [0, 0, 0, 0, 0], sizes = [%11, %12, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%11, %12}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%7 = flow.dispatch.workgroups[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5} =
(%arg2: index, %arg3: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg4: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg5: index, %arg6: index, %arg7: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%10 = flow.dispatch.workload.ordinal %arg5, 1 : index
%11 = flow.dispatch.workload.ordinal %arg6, 2 : index
%12 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%10, %11}
%13 = flow.dispatch.tie_shape %arg7 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%10, %11}
%14 = flow.dispatch.workload.ordinal %arg2, 0 : index
%15 = flow.dispatch.tensor.load %arg3, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%16 = flow.dispatch.tensor.load %12, offsets = [0, 0, 0, 0, 0], sizes = [%10, %11, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%10, %11} -> tensor<?x?x3200x8x1xf32>
%17 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%14]
%18 = tensor.empty(%10) : tensor<?x2160x3200x4x1xf16>
%19 = tensor.empty(%10, %17) : tensor<?x?x2160x8x4xf32>
%20 = linalg.fill ins(%cst : f32) outs(%19 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%21 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%15 : tensor<2160x3200x4x1xf16>) outs(%18 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%22 = linalg.batch_mmt4d ins(%16, %21 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%20 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %22, %13, offsets = [0, 0, 0, 0, 0], sizes = [%10, %11, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%10, %11}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%8 = flow.dispatch.workgroups[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1} =
(%arg2: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg3: index, %arg4: index, %arg5: index, %arg6: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%10 = flow.dispatch.workload.ordinal %arg3, 0 : index
%11 = flow.dispatch.workload.ordinal %arg4, 1 : index
%12 = flow.dispatch.workload.ordinal %arg5, 2 : index
%13 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%11, %10}
%14 = flow.dispatch.tie_shape %arg6 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%11, %12}
%15 = flow.dispatch.tensor.load %13, offsets = [0, 0, 0, 0, 0], sizes = [%11, %10, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%11, %10} -> tensor<?x?x2160x8x4xf32>
%16 = tensor.empty(%11, %12) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %15 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %16 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %14, offsets = [0, 0, 0], sizes = [%11, %12, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%11, %12}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch.workgroups(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16> =
(%arg2: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg3: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%10 = flow.dispatch.tensor.load %arg2, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%11 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %10 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %11 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg3, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
flow.return
} count() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
%5 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%6 = flow.dispatch.workgroups[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5} =
(%arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg3: index, %arg4: index, %arg5: index, %arg6: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%10 = flow.dispatch.workload.ordinal %arg3, 0 : index
%11 = flow.dispatch.workload.ordinal %arg4, 1 : index
%12 = flow.dispatch.workload.ordinal %arg5, 2 : index
%13 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%11, %10}
%14 = flow.dispatch.tie_shape %arg6 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%11, %12}
%15 = flow.dispatch.tensor.load %13, offsets = [0, 0, 0], sizes = [%11, %10, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%11, %10} -> tensor<?x?x3200xf32>
%16 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%10]
%17 = tensor.empty(%11, %16) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %15 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %17 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %14, offsets = [0, 0, 0, 0, 0], sizes = [%11, %12, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%11, %12}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%7 = flow.dispatch.workgroups[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5} =
(%arg2: index, %arg3: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg4: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg5: index, %arg6: index, %arg7: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%10 = flow.dispatch.workload.ordinal %arg5, 1 : index
%11 = flow.dispatch.workload.ordinal %arg6, 2 : index
%12 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%10, %11}
%13 = flow.dispatch.tie_shape %arg7 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%10, %11}
%14 = flow.dispatch.workload.ordinal %arg2, 0 : index
%15 = flow.dispatch.tensor.load %arg3, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%16 = flow.dispatch.tensor.load %12, offsets = [0, 0, 0, 0, 0], sizes = [%10, %11, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%10, %11} -> tensor<?x?x3200x8x1xf32>
%17 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%14]
%18 = tensor.empty(%10) : tensor<?x2160x3200x4x1xf16>
%19 = tensor.empty(%10, %17) : tensor<?x?x2160x8x4xf32>
%20 = linalg.fill ins(%cst : f32) outs(%19 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%21 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%15 : tensor<2160x3200x4x1xf16>) outs(%18 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%22 = linalg.batch_mmt4d ins(%16, %21 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%20 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %22, %13, offsets = [0, 0, 0, 0, 0], sizes = [%10, %11, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%10, %11}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%8 = flow.dispatch.workgroups[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1} =
(%arg2: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg3: index, %arg4: index, %arg5: index, %arg6: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%10 = flow.dispatch.workload.ordinal %arg3, 0 : index
%11 = flow.dispatch.workload.ordinal %arg4, 1 : index
%12 = flow.dispatch.workload.ordinal %arg5, 2 : index
%13 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%11, %10}
%14 = flow.dispatch.tie_shape %arg6 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%11, %12}
%15 = flow.dispatch.tensor.load %13, offsets = [0, 0, 0, 0, 0], sizes = [%11, %10, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%11, %10} -> tensor<?x?x2160x8x4xf32>
%16 = tensor.empty(%11, %12) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %15 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %16 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %14, offsets = [0, 0, 0], sizes = [%11, %12, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%11, %12}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After InitializeEmptyTensorsPass (iree-flow-initialize-empty-tensors) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch.workgroups(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16> =
(%arg2: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg3: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%10 = flow.dispatch.tensor.load %arg2, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%11 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %10 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %11 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg3, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
flow.return
} count() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
%5 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%6 = flow.dispatch.workgroups[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5} =
(%arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg3: index, %arg4: index, %arg5: index, %arg6: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%10 = flow.dispatch.workload.ordinal %arg3, 0 : index
%11 = flow.dispatch.workload.ordinal %arg4, 1 : index
%12 = flow.dispatch.workload.ordinal %arg5, 2 : index
%13 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%11, %10}
%14 = flow.dispatch.tie_shape %arg6 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%11, %12}
%15 = flow.dispatch.tensor.load %13, offsets = [0, 0, 0], sizes = [%11, %10, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%11, %10} -> tensor<?x?x3200xf32>
%16 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%10]
%17 = tensor.empty(%11, %16) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %15 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %17 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %14, offsets = [0, 0, 0, 0, 0], sizes = [%11, %12, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%11, %12}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%7 = flow.dispatch.workgroups[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5} =
(%arg2: index, %arg3: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg4: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg5: index, %arg6: index, %arg7: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%10 = flow.dispatch.workload.ordinal %arg5, 1 : index
%11 = flow.dispatch.workload.ordinal %arg6, 2 : index
%12 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%10, %11}
%13 = flow.dispatch.tie_shape %arg7 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%10, %11}
%14 = flow.dispatch.workload.ordinal %arg2, 0 : index
%15 = flow.dispatch.tensor.load %arg3, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%16 = flow.dispatch.tensor.load %12, offsets = [0, 0, 0, 0, 0], sizes = [%10, %11, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%10, %11} -> tensor<?x?x3200x8x1xf32>
%17 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%14]
%18 = tensor.empty(%10) : tensor<?x2160x3200x4x1xf16>
%19 = tensor.empty(%10, %17) : tensor<?x?x2160x8x4xf32>
%20 = linalg.fill ins(%cst : f32) outs(%19 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%21 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%15 : tensor<2160x3200x4x1xf16>) outs(%18 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%22 = linalg.batch_mmt4d ins(%16, %21 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%20 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %22, %13, offsets = [0, 0, 0, 0, 0], sizes = [%10, %11, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%10, %11}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%8 = flow.dispatch.workgroups[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1} =
(%arg2: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg3: index, %arg4: index, %arg5: index, %arg6: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%10 = flow.dispatch.workload.ordinal %arg3, 0 : index
%11 = flow.dispatch.workload.ordinal %arg4, 1 : index
%12 = flow.dispatch.workload.ordinal %arg5, 2 : index
%13 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%11, %10}
%14 = flow.dispatch.tie_shape %arg6 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%11, %12}
%15 = flow.dispatch.tensor.load %13, offsets = [0, 0, 0, 0, 0], sizes = [%11, %10, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%11, %10} -> tensor<?x?x2160x8x4xf32>
%16 = tensor.empty(%11, %12) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %15 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %16 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %14, offsets = [0, 0, 0], sizes = [%11, %12, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%11, %12}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After OutlineDispatchExternsPass (iree-flow-outline-dispatch-externs) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch.workgroups(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16> =
(%arg2: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg3: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%10 = flow.dispatch.tensor.load %arg2, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%11 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %10 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %11 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg3, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
flow.return
} count() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
%5 = affine.apply #map()[%1]
%6 = flow.dispatch.workgroups[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5} =
(%arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg3: index, %arg4: index, %arg5: index, %arg6: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%10 = flow.dispatch.workload.ordinal %arg3, 0 : index
%11 = flow.dispatch.workload.ordinal %arg4, 1 : index
%12 = flow.dispatch.workload.ordinal %arg5, 2 : index
%13 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%11, %10}
%14 = flow.dispatch.tie_shape %arg6 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%11, %12}
%15 = flow.dispatch.tensor.load %13, offsets = [0, 0, 0], sizes = [%11, %10, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%11, %10} -> tensor<?x?x3200xf32>
%16 = affine.apply #map()[%10]
%17 = tensor.empty(%11, %16) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %15 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %17 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %14, offsets = [0, 0, 0, 0, 0], sizes = [%11, %12, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%11, %12}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%7 = flow.dispatch.workgroups[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5} =
(%arg2: index, %arg3: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg4: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg5: index, %arg6: index, %arg7: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%10 = flow.dispatch.workload.ordinal %arg5, 1 : index
%11 = flow.dispatch.workload.ordinal %arg6, 2 : index
%12 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%10, %11}
%13 = flow.dispatch.tie_shape %arg7 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%10, %11}
%14 = flow.dispatch.workload.ordinal %arg2, 0 : index
%15 = flow.dispatch.tensor.load %arg3, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%16 = flow.dispatch.tensor.load %12, offsets = [0, 0, 0, 0, 0], sizes = [%10, %11, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%10, %11} -> tensor<?x?x3200x8x1xf32>
%17 = affine.apply #map()[%14]
%18 = tensor.empty(%10) : tensor<?x2160x3200x4x1xf16>
%19 = tensor.empty(%10, %17) : tensor<?x?x2160x8x4xf32>
%20 = linalg.fill ins(%cst : f32) outs(%19 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%21 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%15 : tensor<2160x3200x4x1xf16>) outs(%18 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%22 = linalg.batch_mmt4d ins(%16, %21 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%20 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %22, %13, offsets = [0, 0, 0, 0, 0], sizes = [%10, %11, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%10, %11}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%8 = flow.dispatch.workgroups[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1} =
(%arg2: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg3: index, %arg4: index, %arg5: index, %arg6: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%10 = flow.dispatch.workload.ordinal %arg3, 0 : index
%11 = flow.dispatch.workload.ordinal %arg4, 1 : index
%12 = flow.dispatch.workload.ordinal %arg5, 2 : index
%13 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%11, %10}
%14 = flow.dispatch.tie_shape %arg6 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%11, %12}
%15 = flow.dispatch.tensor.load %13, offsets = [0, 0, 0, 0, 0], sizes = [%11, %10, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%11, %10} -> tensor<?x?x2160x8x4xf32>
%16 = tensor.empty(%11, %12) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %15 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %16 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %14, offsets = [0, 0, 0], sizes = [%11, %12, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%11, %12}
flow.return
} count(%arg2: index, %arg3: index, %arg4: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg2, %arg3, %arg4
flow.return %x, %y, %z : index, index, index
}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After OutlineDispatchRegionsPass (iree-flow-outline-dispatch-regions) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
flow.executable private @matmul_broad_dispatch_0 {
flow.executable.export public @matmul_broad_dispatch_0 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0(%arg0: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg1: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%0 = flow.dispatch.tensor.load %arg0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%1 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %1 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
flow.executable private @matmul_broad_dispatch_1 {
flow.executable.export public @matmul_broad_dispatch_1 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
flow.executable private @matmul_broad_dispatch_2 {
flow.executable.export public @matmul_broad_dispatch_2 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2(%arg0: index, %arg1: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg3: index, %arg4: index, %arg5: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg3, 1 : index
%1 = flow.dispatch.workload.ordinal %arg4, 2 : index
%2 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1}
%3 = flow.dispatch.tie_shape %arg5 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
%4 = flow.dispatch.workload.ordinal %arg0, 0 : index
%5 = flow.dispatch.tensor.load %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.tensor.load %2, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1} -> tensor<?x?x3200x8x1xf32>
%7 = affine.apply #map()[%4]
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%5 : tensor<2160x3200x4x1xf16>) outs(%8 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%12 = linalg.batch_mmt4d ins(%6, %11 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %12, %3, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
return
}
}
}
flow.executable private @matmul_broad_dispatch_3 {
flow.executable.export public @matmul_broad_dispatch_3 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply #map()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After AnnotateDispatchesPass (iree-flow-annotate-dispatches) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
flow.executable private @matmul_broad_dispatch_0 {
flow.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg1: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%0 = flow.dispatch.tensor.load %arg0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%1 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %1 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
flow.executable private @matmul_broad_dispatch_1 {
flow.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
flow.executable private @matmul_broad_dispatch_2 {
flow.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg3: index, %arg4: index, %arg5: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg3, 1 : index
%1 = flow.dispatch.workload.ordinal %arg4, 2 : index
%2 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1}
%3 = flow.dispatch.tie_shape %arg5 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
%4 = flow.dispatch.workload.ordinal %arg0, 0 : index
%5 = flow.dispatch.tensor.load %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.tensor.load %2, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1} -> tensor<?x?x3200x8x1xf32>
%7 = affine.apply #map()[%4]
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%5 : tensor<2160x3200x4x1xf16>) outs(%8 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%12 = linalg.batch_mmt4d ins(%6, %11 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %12, %3, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
return
}
}
}
flow.executable private @matmul_broad_dispatch_3 {
flow.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply #map()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After StripDebugOps (iree-util-strip-debug-ops) //----- //
flow.executable private @matmul_broad_dispatch_0 {
flow.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg1: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%0 = flow.dispatch.tensor.load %arg0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%1 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %1 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
// -----// IR Dump After StripDebugOps (iree-util-strip-debug-ops) //----- //
flow.executable private @matmul_broad_dispatch_1 {
flow.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
// -----// IR Dump After StripDebugOps (iree-util-strip-debug-ops) //----- //
flow.executable private @matmul_broad_dispatch_2 {
flow.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg3: index, %arg4: index, %arg5: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg3, 1 : index
%1 = flow.dispatch.workload.ordinal %arg4, 2 : index
%2 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1}
%3 = flow.dispatch.tie_shape %arg5 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
%4 = flow.dispatch.workload.ordinal %arg0, 0 : index
%5 = flow.dispatch.tensor.load %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.tensor.load %2, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1} -> tensor<?x?x3200x8x1xf32>
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%4]
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%5 : tensor<2160x3200x4x1xf16>) outs(%8 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%12 = linalg.batch_mmt4d ins(%6, %11 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %12, %3, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
return
}
}
}
// -----// IR Dump After StripDebugOps (iree-util-strip-debug-ops) //----- //
flow.executable private @matmul_broad_dispatch_3 {
flow.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After DeduplicateExecutablesPass (iree-flow-deduplicate-executables) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
flow.executable private @matmul_broad_dispatch_0 {
flow.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg1: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%0 = flow.dispatch.tensor.load %arg0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%1 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %1 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
flow.executable private @matmul_broad_dispatch_1 {
flow.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
flow.executable private @matmul_broad_dispatch_2 {
flow.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg3: index, %arg4: index, %arg5: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg3, 1 : index
%1 = flow.dispatch.workload.ordinal %arg4, 2 : index
%2 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1}
%3 = flow.dispatch.tie_shape %arg5 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
%4 = flow.dispatch.workload.ordinal %arg0, 0 : index
%5 = flow.dispatch.tensor.load %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.tensor.load %2, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1} -> tensor<?x?x3200x8x1xf32>
%7 = affine.apply #map()[%4]
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%5 : tensor<2160x3200x4x1xf16>) outs(%8 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%12 = linalg.batch_mmt4d ins(%6, %11 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %12, %3, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
return
}
}
}
flow.executable private @matmul_broad_dispatch_3 {
flow.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply #map()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After InjectTensorTracingPass (iree-flow-inject-tensor-tracing) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After CleanupTensorShapesPass (iree-flow-cleanup-tensor-shapes) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After OutlineConstantsPass (iree-flow-outline-constants) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local], iree.fixedpoint.iteration = 0 : index} {
flow.executable private @matmul_broad_dispatch_0 {
flow.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg1: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%0 = flow.dispatch.tensor.load %arg0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%1 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %1 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
flow.executable private @matmul_broad_dispatch_1 {
flow.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
flow.executable private @matmul_broad_dispatch_2 {
flow.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg3: index, %arg4: index, %arg5: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg3, 1 : index
%1 = flow.dispatch.workload.ordinal %arg4, 2 : index
%2 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1}
%3 = flow.dispatch.tie_shape %arg5 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
%4 = flow.dispatch.workload.ordinal %arg0, 0 : index
%5 = flow.dispatch.tensor.load %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.tensor.load %2, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1} -> tensor<?x?x3200x8x1xf32>
%7 = affine.apply #map()[%4]
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%5 : tensor<2160x3200x4x1xf16>) outs(%8 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%12 = linalg.batch_mmt4d ins(%6, %11 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %12, %3, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
return
}
}
}
flow.executable private @matmul_broad_dispatch_3 {
flow.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply #map()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After SimplifyGlobalAccesses (iree-util-simplify-global-accesses) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After ApplyPatterns (iree-util-apply-patterns) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local], iree.fixedpoint.iteration = 0 : index} {
flow.executable private @matmul_broad_dispatch_0 {
flow.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg1: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%0 = flow.dispatch.tensor.load %arg0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%1 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %1 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
flow.executable private @matmul_broad_dispatch_1 {
flow.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
flow.executable private @matmul_broad_dispatch_2 {
flow.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg3: index, %arg4: index, %arg5: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg3, 1 : index
%1 = flow.dispatch.workload.ordinal %arg4, 2 : index
%2 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1}
%3 = flow.dispatch.tie_shape %arg5 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
%4 = flow.dispatch.workload.ordinal %arg0, 0 : index
%5 = flow.dispatch.tensor.load %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.tensor.load %2, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1} -> tensor<?x?x3200x8x1xf32>
%7 = affine.apply #map()[%4]
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%5 : tensor<2160x3200x4x1xf16>) outs(%8 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%12 = linalg.batch_mmt4d ins(%6, %11 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %12, %3, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
return
}
}
}
flow.executable private @matmul_broad_dispatch_3 {
flow.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply #map()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After FoldGlobals (iree-util-fold-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local], iree.fixedpoint.iteration = 0 : index} {
flow.executable private @matmul_broad_dispatch_0 {
flow.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg1: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%0 = flow.dispatch.tensor.load %arg0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%1 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %1 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
flow.executable private @matmul_broad_dispatch_1 {
flow.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
flow.executable private @matmul_broad_dispatch_2 {
flow.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg3: index, %arg4: index, %arg5: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg3, 1 : index
%1 = flow.dispatch.workload.ordinal %arg4, 2 : index
%2 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1}
%3 = flow.dispatch.tie_shape %arg5 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
%4 = flow.dispatch.workload.ordinal %arg0, 0 : index
%5 = flow.dispatch.tensor.load %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.tensor.load %2, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1} -> tensor<?x?x3200x8x1xf32>
%7 = affine.apply #map()[%4]
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%5 : tensor<2160x3200x4x1xf16>) outs(%8 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%12 = linalg.batch_mmt4d ins(%6, %11 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %12, %3, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
return
}
}
}
flow.executable private @matmul_broad_dispatch_3 {
flow.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply #map()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After FuseGlobals (iree-util-fuse-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local], iree.fixedpoint.iteration = 0 : index} {
flow.executable private @matmul_broad_dispatch_0 {
flow.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg1: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%0 = flow.dispatch.tensor.load %arg0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%1 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %1 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
flow.executable private @matmul_broad_dispatch_1 {
flow.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
flow.executable private @matmul_broad_dispatch_2 {
flow.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg3: index, %arg4: index, %arg5: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg3, 1 : index
%1 = flow.dispatch.workload.ordinal %arg4, 2 : index
%2 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1}
%3 = flow.dispatch.tie_shape %arg5 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
%4 = flow.dispatch.workload.ordinal %arg0, 0 : index
%5 = flow.dispatch.tensor.load %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.tensor.load %2, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1} -> tensor<?x?x3200x8x1xf32>
%7 = affine.apply #map()[%4]
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%5 : tensor<2160x3200x4x1xf16>) outs(%8 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%12 = linalg.batch_mmt4d ins(%6, %11 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %12, %3, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
return
}
}
}
flow.executable private @matmul_broad_dispatch_3 {
flow.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply #map()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After IPO (iree-util-ipo) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local], iree.fixedpoint.iteration = 0 : index} {
flow.executable private @matmul_broad_dispatch_0 {
flow.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg1: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%0 = flow.dispatch.tensor.load %arg0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%1 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %1 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
flow.executable private @matmul_broad_dispatch_1 {
flow.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
flow.executable private @matmul_broad_dispatch_2 {
flow.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg3: index, %arg4: index, %arg5: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg3, 1 : index
%1 = flow.dispatch.workload.ordinal %arg4, 2 : index
%2 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1}
%3 = flow.dispatch.tie_shape %arg5 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
%4 = flow.dispatch.workload.ordinal %arg0, 0 : index
%5 = flow.dispatch.tensor.load %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.tensor.load %2, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1} -> tensor<?x?x3200x8x1xf32>
%7 = affine.apply #map()[%4]
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%5 : tensor<2160x3200x4x1xf16>) outs(%8 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%12 = linalg.batch_mmt4d ins(%6, %11 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %12, %3, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
return
}
}
}
flow.executable private @matmul_broad_dispatch_3 {
flow.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply #map()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After FixedPointIterator (iree-util-fixed-point-iterator) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
flow.executable private @matmul_broad_dispatch_0 {
flow.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg1: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%0 = flow.dispatch.tensor.load %arg0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%1 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %1 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
flow.executable private @matmul_broad_dispatch_1 {
flow.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
flow.executable private @matmul_broad_dispatch_2 {
flow.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg3: index, %arg4: index, %arg5: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg3, 1 : index
%1 = flow.dispatch.workload.ordinal %arg4, 2 : index
%2 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1}
%3 = flow.dispatch.tie_shape %arg5 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
%4 = flow.dispatch.workload.ordinal %arg0, 0 : index
%5 = flow.dispatch.tensor.load %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.tensor.load %2, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1} -> tensor<?x?x3200x8x1xf32>
%7 = affine.apply #map()[%4]
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%5 : tensor<2160x3200x4x1xf16>) outs(%8 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%12 = linalg.batch_mmt4d ins(%6, %11 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %12, %3, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
return
}
}
}
flow.executable private @matmul_broad_dispatch_3 {
flow.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply #map()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After SymbolDCE (symbol-dce) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
flow.executable private @matmul_broad_dispatch_0 {
flow.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg1: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%0 = flow.dispatch.tensor.load %arg0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%1 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %1 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
flow.executable private @matmul_broad_dispatch_1 {
flow.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
flow.executable private @matmul_broad_dispatch_2 {
flow.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg3: index, %arg4: index, %arg5: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg3, 1 : index
%1 = flow.dispatch.workload.ordinal %arg4, 2 : index
%2 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1}
%3 = flow.dispatch.tie_shape %arg5 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
%4 = flow.dispatch.workload.ordinal %arg0, 0 : index
%5 = flow.dispatch.tensor.load %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.tensor.load %2, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1} -> tensor<?x?x3200x8x1xf32>
%7 = affine.apply #map()[%4]
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%5 : tensor<2160x3200x4x1xf16>) outs(%8 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%12 = linalg.batch_mmt4d ins(%6, %11 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %12, %3, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
return
}
}
}
flow.executable private @matmul_broad_dispatch_3 {
flow.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply #map()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After VerifyInputPass (iree-stream-verify-input) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
flow.executable private @matmul_broad_dispatch_0 {
flow.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg1: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%0 = flow.dispatch.tensor.load %arg0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%1 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %1 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
flow.executable private @matmul_broad_dispatch_1 {
flow.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
flow.executable private @matmul_broad_dispatch_2 {
flow.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg3: index, %arg4: index, %arg5: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg3, 1 : index
%1 = flow.dispatch.workload.ordinal %arg4, 2 : index
%2 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1}
%3 = flow.dispatch.tie_shape %arg5 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
%4 = flow.dispatch.workload.ordinal %arg0, 0 : index
%5 = flow.dispatch.tensor.load %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.tensor.load %2, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1} -> tensor<?x?x3200x8x1xf32>
%7 = affine.apply #map()[%4]
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%5 : tensor<2160x3200x4x1xf16>) outs(%8 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%12 = linalg.batch_mmt4d ins(%6, %11 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %12, %3, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
return
}
}
}
flow.executable private @matmul_broad_dispatch_3 {
flow.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply #map()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After SimplifyGlobalAccesses (iree-util-simplify-global-accesses) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
// -----// IR Dump After ApplyPatterns (iree-util-apply-patterns) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
flow.executable private @matmul_broad_dispatch_0 {
flow.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg1: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%0 = flow.dispatch.tensor.load %arg0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%1 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %1 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
flow.executable private @matmul_broad_dispatch_1 {
flow.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
flow.executable private @matmul_broad_dispatch_2 {
flow.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg3: index, %arg4: index, %arg5: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg3, 1 : index
%1 = flow.dispatch.workload.ordinal %arg4, 2 : index
%2 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1}
%3 = flow.dispatch.tie_shape %arg5 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
%4 = flow.dispatch.workload.ordinal %arg0, 0 : index
%5 = flow.dispatch.tensor.load %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.tensor.load %2, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1} -> tensor<?x?x3200x8x1xf32>
%7 = affine.apply #map()[%4]
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%5 : tensor<2160x3200x4x1xf16>) outs(%8 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%12 = linalg.batch_mmt4d ins(%6, %11 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %12, %3, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
return
}
}
}
flow.executable private @matmul_broad_dispatch_3 {
flow.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply #map()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After FoldGlobals (iree-util-fold-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
flow.executable private @matmul_broad_dispatch_0 {
flow.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg1: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%0 = flow.dispatch.tensor.load %arg0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%1 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %1 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
flow.executable private @matmul_broad_dispatch_1 {
flow.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
flow.executable private @matmul_broad_dispatch_2 {
flow.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg3: index, %arg4: index, %arg5: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg3, 1 : index
%1 = flow.dispatch.workload.ordinal %arg4, 2 : index
%2 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1}
%3 = flow.dispatch.tie_shape %arg5 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
%4 = flow.dispatch.workload.ordinal %arg0, 0 : index
%5 = flow.dispatch.tensor.load %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.tensor.load %2, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1} -> tensor<?x?x3200x8x1xf32>
%7 = affine.apply #map()[%4]
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%5 : tensor<2160x3200x4x1xf16>) outs(%8 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%12 = linalg.batch_mmt4d ins(%6, %11 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %12, %3, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
return
}
}
}
flow.executable private @matmul_broad_dispatch_3 {
flow.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply #map()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After FuseGlobals (iree-util-fuse-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
flow.executable private @matmul_broad_dispatch_0 {
flow.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg1: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%0 = flow.dispatch.tensor.load %arg0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%1 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %1 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
flow.executable private @matmul_broad_dispatch_1 {
flow.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
flow.executable private @matmul_broad_dispatch_2 {
flow.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg3: index, %arg4: index, %arg5: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg3, 1 : index
%1 = flow.dispatch.workload.ordinal %arg4, 2 : index
%2 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1}
%3 = flow.dispatch.tie_shape %arg5 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
%4 = flow.dispatch.workload.ordinal %arg0, 0 : index
%5 = flow.dispatch.tensor.load %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.tensor.load %2, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1} -> tensor<?x?x3200x8x1xf32>
%7 = affine.apply #map()[%4]
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%5 : tensor<2160x3200x4x1xf16>) outs(%8 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%12 = linalg.batch_mmt4d ins(%6, %11 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %12, %3, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
return
}
}
}
flow.executable private @matmul_broad_dispatch_3 {
flow.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply #map()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After IPO (iree-util-ipo) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
flow.executable private @matmul_broad_dispatch_0 {
flow.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>, %arg1: !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>) {
%0 = flow.dispatch.tensor.load %arg0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%1 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %1 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
flow.executable private @matmul_broad_dispatch_1 {
flow.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
flow.executable private @matmul_broad_dispatch_2 {
flow.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>, %arg2: !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>, %arg3: index, %arg4: index, %arg5: !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>) {
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg3, 1 : index
%1 = flow.dispatch.workload.ordinal %arg4, 2 : index
%2 = flow.dispatch.tie_shape %arg2 : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1}
%3 = flow.dispatch.tie_shape %arg5 : !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
%4 = flow.dispatch.workload.ordinal %arg0, 0 : index
%5 = flow.dispatch.tensor.load %arg1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%6 = flow.dispatch.tensor.load %2, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%0, %1} -> tensor<?x?x3200x8x1xf32>
%7 = affine.apply #map()[%4]
%8 = tensor.empty(%0) : tensor<?x2160x3200x4x1xf16>
%9 = tensor.empty(%0, %7) : tensor<?x?x2160x8x4xf32>
%10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%11 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%5 : tensor<2160x3200x4x1xf16>) outs(%8 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%12 = linalg.batch_mmt4d ins(%6, %11 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %12, %3, offsets = [0, 0, 0, 0, 0], sizes = [%0, %1, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%0, %1}
return
}
}
}
flow.executable private @matmul_broad_dispatch_3 {
flow.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
flow.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>, %arg1: index, %arg2: index, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>) {
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = flow.dispatch.tie_shape %arg0 : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = flow.dispatch.tie_shape %arg4 : !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%2 = hal.tensor.import %arg0 "input0" : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1}
%3 = hal.tensor.import %arg1 "input1" : !hal.buffer_view -> tensor<8640x3200xf16>
%4 = flow.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%3) : (tensor<8640x3200xf16>) -> tensor<2160x3200x4x1xf16>
%5 = affine.apply #map()[%1]
%6 = flow.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %5](%2, %1, %0, %5) : (tensor<?x?x3200xf32>{%0, %1}, index, index, index) -> tensor<?x?x3200x8x1xf32>{%0, %5}
%7 = flow.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %5](%1, %4, %6, %0, %5) : (index, tensor<2160x3200x4x1xf16>, tensor<?x?x3200x8x1xf32>{%0, %5}, index, index) -> tensor<?x?x2160x8x4xf32>{%0, %5}
%8 = flow.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%5, %0, %1](%7, %5, %0, %1) : (tensor<?x?x2160x8x4xf32>{%0, %5}, index, index, index) -> tensor<?x?x8640xf32>{%0, %1}
%9 = hal.tensor.export %8 "output0" : tensor<?x?x8640xf32>{%0, %1} -> !hal.buffer_view
util.return %9 : !hal.buffer_view
}
}
// -----// IR Dump After ConvertToStreamPass (iree-stream-conversion) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%cst = arith.constant 0.000000e+00 : f32
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
%c3200 = arith.constant 3200 : index
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = stream.tensor.sizeof tensor<?x?x3200xf32>{%0, %1} : index
%3 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%2}
%4 = stream.async.transfer %3 : !stream.resource<external>{%2} -> !stream.resource<*>{%2}
%element_type_f16 = hal.element_type<f16> : i32
%dense_row_major_0 = hal.encoding_type<dense_row_major> : i32
%c8640 = arith.constant 8640 : index
%c3200_1 = arith.constant 3200 : index
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200_1]) type(%element_type_f16) encoding(%dense_row_major_0)
%5 = stream.tensor.sizeof tensor<8640x3200xf16> : index
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%5}
%7 = stream.async.transfer %6 : !stream.resource<external>{%5} -> !stream.resource<*>{%5}
%c0 = arith.constant 0 : index
%8 = stream.tensor.sizeof tensor<2160x3200x4x1xf16> : index
%9 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %5 for %5]) : (!stream.resource<*>{%5}) -> !stream.resource<*>{%8}
%10 = affine.apply #map()[%1]
%c0_2 = arith.constant 0 : index
%11 = stream.tensor.sizeof tensor<?x?x3200x8x1xf32>{%0, %10} : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %10](%4[%c0_2 to %2 for %2], %1, %0, %10) : (!stream.resource<*>{%2}, index, index, index) -> !stream.resource<*>{%11}
%c0_3 = arith.constant 0 : index
%13 = stream.tensor.sizeof tensor<?x?x2160x8x4xf32>{%0, %10} : index
%14 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %10](%1, %9[%c0_3 to %8 for %8], %12[%c0_3 to %11 for %11], %0, %10) : (index, !stream.resource<*>{%8}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%13}
%c0_4 = arith.constant 0 : index
%15 = stream.tensor.sizeof tensor<?x?x8640xf32>{%0, %1} : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%10, %0, %1](%14[%c0_4 to %13 for %13], %10, %0, %1) : (!stream.resource<*>{%13}, index, index, index) -> !stream.resource<*>{%15}
%17 = stream.async.transfer %16 : !stream.resource<*>{%15} -> !stream.resource<external>{%15}
%18 = stream.tensor.export %17 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %18 : !hal.buffer_view
}
}
// -----// IR Dump After VerifyLoweringToTensorsPass (iree-stream-verify-lowering-to-tensors) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%cst = arith.constant 0.000000e+00 : f32
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
%c3200 = arith.constant 3200 : index
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = stream.tensor.sizeof tensor<?x?x3200xf32>{%0, %1} : index
%3 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%2}
%4 = stream.async.transfer %3 : !stream.resource<external>{%2} -> !stream.resource<*>{%2}
%element_type_f16 = hal.element_type<f16> : i32
%dense_row_major_0 = hal.encoding_type<dense_row_major> : i32
%c8640 = arith.constant 8640 : index
%c3200_1 = arith.constant 3200 : index
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200_1]) type(%element_type_f16) encoding(%dense_row_major_0)
%5 = stream.tensor.sizeof tensor<8640x3200xf16> : index
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%5}
%7 = stream.async.transfer %6 : !stream.resource<external>{%5} -> !stream.resource<*>{%5}
%c0 = arith.constant 0 : index
%8 = stream.tensor.sizeof tensor<2160x3200x4x1xf16> : index
%9 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %5 for %5]) : (!stream.resource<*>{%5}) -> !stream.resource<*>{%8}
%10 = affine.apply #map()[%1]
%c0_2 = arith.constant 0 : index
%11 = stream.tensor.sizeof tensor<?x?x3200x8x1xf32>{%0, %10} : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %10](%4[%c0_2 to %2 for %2], %1, %0, %10) : (!stream.resource<*>{%2}, index, index, index) -> !stream.resource<*>{%11}
%c0_3 = arith.constant 0 : index
%13 = stream.tensor.sizeof tensor<?x?x2160x8x4xf32>{%0, %10} : index
%14 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %10](%1, %9[%c0_3 to %8 for %8], %12[%c0_3 to %11 for %11], %0, %10) : (index, !stream.resource<*>{%8}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%13}
%c0_4 = arith.constant 0 : index
%15 = stream.tensor.sizeof tensor<?x?x8640xf32>{%0, %1} : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%10, %0, %1](%14[%c0_4 to %13 for %13], %10, %0, %1) : (!stream.resource<*>{%13}, index, index, index) -> !stream.resource<*>{%15}
%17 = stream.async.transfer %16 : !stream.resource<*>{%15} -> !stream.resource<external>{%15}
%18 = stream.tensor.export %17 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %18 : !hal.buffer_view
}
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = stream.tensor.sizeof tensor<?x?x3200xf32>{%0, %1} : index
%3 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%2}
%4 = stream.async.transfer %3 : !stream.resource<external>{%2} -> !stream.resource<*>{%2}
%element_type_f16 = hal.element_type<f16> : i32
%dense_row_major_0 = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major_0)
%5 = stream.tensor.sizeof tensor<8640x3200xf16> : index
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%5}
%7 = stream.async.transfer %6 : !stream.resource<external>{%5} -> !stream.resource<*>{%5}
%8 = stream.tensor.sizeof tensor<2160x3200x4x1xf16> : index
%9 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %5 for %5]) : (!stream.resource<*>{%5}) -> !stream.resource<*>{%8}
%10 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%11 = stream.tensor.sizeof tensor<?x?x3200x8x1xf32>{%0, %10} : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %10](%4[%c0 to %2 for %2], %1, %0, %10) : (!stream.resource<*>{%2}, index, index, index) -> !stream.resource<*>{%11}
%13 = stream.tensor.sizeof tensor<?x?x2160x8x4xf32>{%0, %10} : index
%14 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %10](%1, %9[%c0 to %8 for %8], %12[%c0 to %11 for %11], %0, %10) : (index, !stream.resource<*>{%8}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%13}
%15 = stream.tensor.sizeof tensor<?x?x8640xf32>{%0, %1} : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%10, %0, %1](%14[%c0 to %13 for %13], %10, %0, %1) : (!stream.resource<*>{%13}, index, index, index) -> !stream.resource<*>{%15}
%17 = stream.async.transfer %16 : !stream.resource<*>{%15} -> !stream.resource<external>{%15}
%18 = stream.tensor.export %17 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %18 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = stream.tensor.sizeof tensor<?x?x3200xf32>{%0, %1} : index
%3 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%2}
%4 = stream.async.transfer %3 : !stream.resource<external>{%2} -> !stream.resource<*>{%2}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.sizeof tensor<8640x3200xf16> : index
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%5}
%7 = stream.async.transfer %6 : !stream.resource<external>{%5} -> !stream.resource<*>{%5}
%8 = stream.tensor.sizeof tensor<2160x3200x4x1xf16> : index
%9 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %5 for %5]) : (!stream.resource<*>{%5}) -> !stream.resource<*>{%8}
%10 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%11 = stream.tensor.sizeof tensor<?x?x3200x8x1xf32>{%0, %10} : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %10](%4[%c0 to %2 for %2], %1, %0, %10) : (!stream.resource<*>{%2}, index, index, index) -> !stream.resource<*>{%11}
%13 = stream.tensor.sizeof tensor<?x?x2160x8x4xf32>{%0, %10} : index
%14 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %10](%1, %9[%c0 to %8 for %8], %12[%c0 to %11 for %11], %0, %10) : (index, !stream.resource<*>{%8}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%13}
%15 = stream.tensor.sizeof tensor<?x?x8640xf32>{%0, %1} : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%10, %0, %1](%14[%c0 to %13 for %13], %10, %0, %1) : (!stream.resource<*>{%13}, index, index, index) -> !stream.resource<*>{%15}
%17 = stream.async.transfer %16 : !stream.resource<*>{%15} -> !stream.resource<external>{%15}
%18 = stream.tensor.export %17 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %18 : !hal.buffer_view
}
// -----// IR Dump After SimplifyGlobalAccesses (iree-util-simplify-global-accesses) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = stream.tensor.sizeof tensor<?x?x3200xf32>{%0, %1} : index
%3 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%2}
%4 = stream.async.transfer %3 : !stream.resource<external>{%2} -> !stream.resource<*>{%2}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.sizeof tensor<8640x3200xf16> : index
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%5}
%7 = stream.async.transfer %6 : !stream.resource<external>{%5} -> !stream.resource<*>{%5}
%8 = stream.tensor.sizeof tensor<2160x3200x4x1xf16> : index
%9 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %5 for %5]) : (!stream.resource<*>{%5}) -> !stream.resource<*>{%8}
%10 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%11 = stream.tensor.sizeof tensor<?x?x3200x8x1xf32>{%0, %10} : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %10](%4[%c0 to %2 for %2], %1, %0, %10) : (!stream.resource<*>{%2}, index, index, index) -> !stream.resource<*>{%11}
%13 = stream.tensor.sizeof tensor<?x?x2160x8x4xf32>{%0, %10} : index
%14 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %10](%1, %9[%c0 to %8 for %8], %12[%c0 to %11 for %11], %0, %10) : (index, !stream.resource<*>{%8}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%13}
%15 = stream.tensor.sizeof tensor<?x?x8640xf32>{%0, %1} : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%10, %0, %1](%14[%c0 to %13 for %13], %10, %0, %1) : (!stream.resource<*>{%13}, index, index, index) -> !stream.resource<*>{%15}
%17 = stream.async.transfer %16 : !stream.resource<*>{%15} -> !stream.resource<external>{%15}
%18 = stream.tensor.export %17 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %18 : !hal.buffer_view
}
// -----// IR Dump After ApplyPatterns (iree-util-apply-patterns) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = stream.tensor.sizeof tensor<?x?x3200xf32>{%0, %1} : index
%3 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%2}
%4 = stream.async.transfer %3 : !stream.resource<external>{%2} -> !stream.resource<*>{%2}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.sizeof tensor<8640x3200xf16> : index
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%5}
%7 = stream.async.transfer %6 : !stream.resource<external>{%5} -> !stream.resource<*>{%5}
%8 = stream.tensor.sizeof tensor<2160x3200x4x1xf16> : index
%9 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %5 for %5]) : (!stream.resource<*>{%5}) -> !stream.resource<*>{%8}
%10 = affine.apply #map()[%1]
%11 = stream.tensor.sizeof tensor<?x?x3200x8x1xf32>{%0, %10} : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %10](%4[%c0 to %2 for %2], %1, %0, %10) : (!stream.resource<*>{%2}, index, index, index) -> !stream.resource<*>{%11}
%13 = stream.tensor.sizeof tensor<?x?x2160x8x4xf32>{%0, %10} : index
%14 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %10](%1, %9[%c0 to %8 for %8], %12[%c0 to %11 for %11], %0, %10) : (index, !stream.resource<*>{%8}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%13}
%15 = stream.tensor.sizeof tensor<?x?x8640xf32>{%0, %1} : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%10, %0, %1](%14[%c0 to %13 for %13], %10, %0, %1) : (!stream.resource<*>{%13}, index, index, index) -> !stream.resource<*>{%15}
%17 = stream.async.transfer %16 : !stream.resource<*>{%15} -> !stream.resource<external>{%15}
%18 = stream.tensor.export %17 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %18 : !hal.buffer_view
}
}
// -----// IR Dump After FoldGlobals (iree-util-fold-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = stream.tensor.sizeof tensor<?x?x3200xf32>{%0, %1} : index
%3 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%2}
%4 = stream.async.transfer %3 : !stream.resource<external>{%2} -> !stream.resource<*>{%2}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.sizeof tensor<8640x3200xf16> : index
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%5}
%7 = stream.async.transfer %6 : !stream.resource<external>{%5} -> !stream.resource<*>{%5}
%8 = stream.tensor.sizeof tensor<2160x3200x4x1xf16> : index
%9 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %5 for %5]) : (!stream.resource<*>{%5}) -> !stream.resource<*>{%8}
%10 = affine.apply #map()[%1]
%11 = stream.tensor.sizeof tensor<?x?x3200x8x1xf32>{%0, %10} : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %10](%4[%c0 to %2 for %2], %1, %0, %10) : (!stream.resource<*>{%2}, index, index, index) -> !stream.resource<*>{%11}
%13 = stream.tensor.sizeof tensor<?x?x2160x8x4xf32>{%0, %10} : index
%14 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %10](%1, %9[%c0 to %8 for %8], %12[%c0 to %11 for %11], %0, %10) : (index, !stream.resource<*>{%8}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%13}
%15 = stream.tensor.sizeof tensor<?x?x8640xf32>{%0, %1} : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%10, %0, %1](%14[%c0 to %13 for %13], %10, %0, %1) : (!stream.resource<*>{%13}, index, index, index) -> !stream.resource<*>{%15}
%17 = stream.async.transfer %16 : !stream.resource<*>{%15} -> !stream.resource<external>{%15}
%18 = stream.tensor.export %17 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %18 : !hal.buffer_view
}
}
// -----// IR Dump After FuseGlobals (iree-util-fuse-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = stream.tensor.sizeof tensor<?x?x3200xf32>{%0, %1} : index
%3 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%2}
%4 = stream.async.transfer %3 : !stream.resource<external>{%2} -> !stream.resource<*>{%2}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.sizeof tensor<8640x3200xf16> : index
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%5}
%7 = stream.async.transfer %6 : !stream.resource<external>{%5} -> !stream.resource<*>{%5}
%8 = stream.tensor.sizeof tensor<2160x3200x4x1xf16> : index
%9 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %5 for %5]) : (!stream.resource<*>{%5}) -> !stream.resource<*>{%8}
%10 = affine.apply #map()[%1]
%11 = stream.tensor.sizeof tensor<?x?x3200x8x1xf32>{%0, %10} : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %10](%4[%c0 to %2 for %2], %1, %0, %10) : (!stream.resource<*>{%2}, index, index, index) -> !stream.resource<*>{%11}
%13 = stream.tensor.sizeof tensor<?x?x2160x8x4xf32>{%0, %10} : index
%14 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %10](%1, %9[%c0 to %8 for %8], %12[%c0 to %11 for %11], %0, %10) : (index, !stream.resource<*>{%8}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%13}
%15 = stream.tensor.sizeof tensor<?x?x8640xf32>{%0, %1} : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%10, %0, %1](%14[%c0 to %13 for %13], %10, %0, %1) : (!stream.resource<*>{%13}, index, index, index) -> !stream.resource<*>{%15}
%17 = stream.async.transfer %16 : !stream.resource<*>{%15} -> !stream.resource<external>{%15}
%18 = stream.tensor.export %17 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %18 : !hal.buffer_view
}
}
// -----// IR Dump After IPO (iree-util-ipo) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = stream.tensor.sizeof tensor<?x?x3200xf32>{%0, %1} : index
%3 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%2}
%4 = stream.async.transfer %3 : !stream.resource<external>{%2} -> !stream.resource<*>{%2}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.sizeof tensor<8640x3200xf16> : index
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%5}
%7 = stream.async.transfer %6 : !stream.resource<external>{%5} -> !stream.resource<*>{%5}
%8 = stream.tensor.sizeof tensor<2160x3200x4x1xf16> : index
%9 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %5 for %5]) : (!stream.resource<*>{%5}) -> !stream.resource<*>{%8}
%10 = affine.apply #map()[%1]
%11 = stream.tensor.sizeof tensor<?x?x3200x8x1xf32>{%0, %10} : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %10](%4[%c0 to %2 for %2], %1, %0, %10) : (!stream.resource<*>{%2}, index, index, index) -> !stream.resource<*>{%11}
%13 = stream.tensor.sizeof tensor<?x?x2160x8x4xf32>{%0, %10} : index
%14 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %10](%1, %9[%c0 to %8 for %8], %12[%c0 to %11 for %11], %0, %10) : (index, !stream.resource<*>{%8}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%13}
%15 = stream.tensor.sizeof tensor<?x?x8640xf32>{%0, %1} : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%10, %0, %1](%14[%c0 to %13 for %13], %10, %0, %1) : (!stream.resource<*>{%13}, index, index, index) -> !stream.resource<*>{%15}
%17 = stream.async.transfer %16 : !stream.resource<*>{%15} -> !stream.resource<external>{%15}
%18 = stream.tensor.export %17 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %18 : !hal.buffer_view
}
}
// -----// IR Dump After CombineInitializers (iree-util-combine-initializers) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = stream.tensor.sizeof tensor<?x?x3200xf32>{%0, %1} : index
%3 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%2}
%4 = stream.async.transfer %3 : !stream.resource<external>{%2} -> !stream.resource<*>{%2}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.sizeof tensor<8640x3200xf16> : index
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%5}
%7 = stream.async.transfer %6 : !stream.resource<external>{%5} -> !stream.resource<*>{%5}
%8 = stream.tensor.sizeof tensor<2160x3200x4x1xf16> : index
%9 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %5 for %5]) : (!stream.resource<*>{%5}) -> !stream.resource<*>{%8}
%10 = affine.apply #map()[%1]
%11 = stream.tensor.sizeof tensor<?x?x3200x8x1xf32>{%0, %10} : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %10](%4[%c0 to %2 for %2], %1, %0, %10) : (!stream.resource<*>{%2}, index, index, index) -> !stream.resource<*>{%11}
%13 = stream.tensor.sizeof tensor<?x?x2160x8x4xf32>{%0, %10} : index
%14 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %10](%1, %9[%c0 to %8 for %8], %12[%c0 to %11 for %11], %0, %10) : (index, !stream.resource<*>{%8}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%13}
%15 = stream.tensor.sizeof tensor<?x?x8640xf32>{%0, %1} : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%10, %0, %1](%14[%c0 to %13 for %13], %10, %0, %1) : (!stream.resource<*>{%13}, index, index, index) -> !stream.resource<*>{%15}
%17 = stream.async.transfer %16 : !stream.resource<*>{%15} -> !stream.resource<external>{%15}
%18 = stream.tensor.export %17 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %18 : !hal.buffer_view
}
}
// -----// IR Dump After EncodeDeviceTensorsPass (iree-stream-encode-device-tensors) //----- //
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
// -----// IR Dump After EncodeDeviceTensorsPass (iree-stream-encode-device-tensors) //----- //
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
// -----// IR Dump After EncodeDeviceTensorsPass (iree-stream-encode-device-tensors) //----- //
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>, affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
// -----// IR Dump After EncodeDeviceTensorsPass (iree-stream-encode-device-tensors) //----- //
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
// -----// IR Dump After EncodeHostTensorsPass (iree-stream-encode-host-tensors) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%5 = stream.async.transfer %4 : !stream.resource<external>{%3} -> !stream.resource<*>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%7 = stream.async.transfer %6 : !stream.resource<external>{%c55296000} -> !stream.resource<*>{%c55296000}
%8 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %c55296000 for %c55296000]) : (!stream.resource<*>{%c55296000}) -> !stream.resource<*>{%c55296000}
%9 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%10 = arith.muli %0, %c102400 : index
%11 = arith.muli %10, %9 : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %9](%5[%c0 to %3 for %3], %1, %0, %9) : (!stream.resource<*>{%3}, index, index, index) -> !stream.resource<*>{%11}
%13 = arith.muli %0, %c276480 : index
%14 = arith.muli %13, %9 : index
%15 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %9](%1, %8[%c0 to %c55296000 for %c55296000], %12[%c0 to %11 for %11], %0, %9) : (index, !stream.resource<*>{%c55296000}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%14}
%16 = arith.muli %0, %c34560 : index
%17 = arith.muli %16, %1 : index
%18 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%9, %0, %1](%15[%c0 to %14 for %14], %9, %0, %1) : (!stream.resource<*>{%14}, index, index, index) -> !stream.resource<*>{%17}
%19 = stream.async.transfer %18 : !stream.resource<*>{%17} -> !stream.resource<external>{%17}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%17} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%5 = stream.async.transfer %4 : !stream.resource<external>{%3} -> !stream.resource<*>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%7 = stream.async.transfer %6 : !stream.resource<external>{%c55296000} -> !stream.resource<*>{%c55296000}
%8 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %c55296000 for %c55296000]) : (!stream.resource<*>{%c55296000}) -> !stream.resource<*>{%c55296000}
%9 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%10 = arith.muli %0, %c102400 : index
%11 = arith.muli %10, %9 : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %9](%5[%c0 to %3 for %3], %1, %0, %9) : (!stream.resource<*>{%3}, index, index, index) -> !stream.resource<*>{%11}
%13 = arith.muli %0, %c276480 : index
%14 = arith.muli %13, %9 : index
%15 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %9](%1, %8[%c0 to %c55296000 for %c55296000], %12[%c0 to %11 for %11], %0, %9) : (index, !stream.resource<*>{%c55296000}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%14}
%16 = arith.muli %0, %c34560 : index
%17 = arith.muli %16, %1 : index
%18 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%9, %0, %1](%15[%c0 to %14 for %14], %9, %0, %1) : (!stream.resource<*>{%14}, index, index, index) -> !stream.resource<*>{%17}
%19 = stream.async.transfer %18 : !stream.resource<*>{%17} -> !stream.resource<external>{%17}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%17} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%5 = stream.async.transfer %4 : !stream.resource<external>{%3} -> !stream.resource<*>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%7 = stream.async.transfer %6 : !stream.resource<external>{%c55296000} -> !stream.resource<*>{%c55296000}
%8 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %c55296000 for %c55296000]) : (!stream.resource<*>{%c55296000}) -> !stream.resource<*>{%c55296000}
%9 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%10 = arith.muli %0, %c102400 : index
%11 = arith.muli %10, %9 : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %9](%5[%c0 to %3 for %3], %1, %0, %9) : (!stream.resource<*>{%3}, index, index, index) -> !stream.resource<*>{%11}
%13 = arith.muli %0, %c276480 : index
%14 = arith.muli %13, %9 : index
%15 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %9](%1, %8[%c0 to %c55296000 for %c55296000], %12[%c0 to %11 for %11], %0, %9) : (index, !stream.resource<*>{%c55296000}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%14}
%16 = arith.muli %0, %c34560 : index
%17 = arith.muli %16, %1 : index
%18 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%9, %0, %1](%15[%c0 to %14 for %14], %9, %0, %1) : (!stream.resource<*>{%14}, index, index, index) -> !stream.resource<*>{%17}
%19 = stream.async.transfer %18 : !stream.resource<*>{%17} -> !stream.resource<external>{%17}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%17} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
// -----// IR Dump After SimplifyGlobalAccesses (iree-util-simplify-global-accesses) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%5 = stream.async.transfer %4 : !stream.resource<external>{%3} -> !stream.resource<*>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%7 = stream.async.transfer %6 : !stream.resource<external>{%c55296000} -> !stream.resource<*>{%c55296000}
%8 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %c55296000 for %c55296000]) : (!stream.resource<*>{%c55296000}) -> !stream.resource<*>{%c55296000}
%9 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%10 = arith.muli %0, %c102400 : index
%11 = arith.muli %10, %9 : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %9](%5[%c0 to %3 for %3], %1, %0, %9) : (!stream.resource<*>{%3}, index, index, index) -> !stream.resource<*>{%11}
%13 = arith.muli %0, %c276480 : index
%14 = arith.muli %13, %9 : index
%15 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %9](%1, %8[%c0 to %c55296000 for %c55296000], %12[%c0 to %11 for %11], %0, %9) : (index, !stream.resource<*>{%c55296000}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%14}
%16 = arith.muli %0, %c34560 : index
%17 = arith.muli %16, %1 : index
%18 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%9, %0, %1](%15[%c0 to %14 for %14], %9, %0, %1) : (!stream.resource<*>{%14}, index, index, index) -> !stream.resource<*>{%17}
%19 = stream.async.transfer %18 : !stream.resource<*>{%17} -> !stream.resource<external>{%17}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%17} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
// -----// IR Dump After ApplyPatterns (iree-util-apply-patterns) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%5 = stream.async.transfer %4 : !stream.resource<external>{%3} -> !stream.resource<*>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%7 = stream.async.transfer %6 : !stream.resource<external>{%c55296000} -> !stream.resource<*>{%c55296000}
%8 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %c55296000 for %c55296000]) : (!stream.resource<*>{%c55296000}) -> !stream.resource<*>{%c55296000}
%9 = affine.apply #map()[%1]
%10 = arith.muli %0, %c102400 : index
%11 = arith.muli %10, %9 : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %9](%5[%c0 to %3 for %3], %1, %0, %9) : (!stream.resource<*>{%3}, index, index, index) -> !stream.resource<*>{%11}
%13 = arith.muli %0, %c276480 : index
%14 = arith.muli %13, %9 : index
%15 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %9](%1, %8[%c0 to %c55296000 for %c55296000], %12[%c0 to %11 for %11], %0, %9) : (index, !stream.resource<*>{%c55296000}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%14}
%16 = arith.muli %0, %c34560 : index
%17 = arith.muli %16, %1 : index
%18 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%9, %0, %1](%15[%c0 to %14 for %14], %9, %0, %1) : (!stream.resource<*>{%14}, index, index, index) -> !stream.resource<*>{%17}
%19 = stream.async.transfer %18 : !stream.resource<*>{%17} -> !stream.resource<external>{%17}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%17} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
}
// -----// IR Dump After FoldGlobals (iree-util-fold-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%5 = stream.async.transfer %4 : !stream.resource<external>{%3} -> !stream.resource<*>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%7 = stream.async.transfer %6 : !stream.resource<external>{%c55296000} -> !stream.resource<*>{%c55296000}
%8 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %c55296000 for %c55296000]) : (!stream.resource<*>{%c55296000}) -> !stream.resource<*>{%c55296000}
%9 = affine.apply #map()[%1]
%10 = arith.muli %0, %c102400 : index
%11 = arith.muli %10, %9 : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %9](%5[%c0 to %3 for %3], %1, %0, %9) : (!stream.resource<*>{%3}, index, index, index) -> !stream.resource<*>{%11}
%13 = arith.muli %0, %c276480 : index
%14 = arith.muli %13, %9 : index
%15 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %9](%1, %8[%c0 to %c55296000 for %c55296000], %12[%c0 to %11 for %11], %0, %9) : (index, !stream.resource<*>{%c55296000}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%14}
%16 = arith.muli %0, %c34560 : index
%17 = arith.muli %16, %1 : index
%18 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%9, %0, %1](%15[%c0 to %14 for %14], %9, %0, %1) : (!stream.resource<*>{%14}, index, index, index) -> !stream.resource<*>{%17}
%19 = stream.async.transfer %18 : !stream.resource<*>{%17} -> !stream.resource<external>{%17}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%17} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
}
// -----// IR Dump After FuseGlobals (iree-util-fuse-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%5 = stream.async.transfer %4 : !stream.resource<external>{%3} -> !stream.resource<*>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%7 = stream.async.transfer %6 : !stream.resource<external>{%c55296000} -> !stream.resource<*>{%c55296000}
%8 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %c55296000 for %c55296000]) : (!stream.resource<*>{%c55296000}) -> !stream.resource<*>{%c55296000}
%9 = affine.apply #map()[%1]
%10 = arith.muli %0, %c102400 : index
%11 = arith.muli %10, %9 : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %9](%5[%c0 to %3 for %3], %1, %0, %9) : (!stream.resource<*>{%3}, index, index, index) -> !stream.resource<*>{%11}
%13 = arith.muli %0, %c276480 : index
%14 = arith.muli %13, %9 : index
%15 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %9](%1, %8[%c0 to %c55296000 for %c55296000], %12[%c0 to %11 for %11], %0, %9) : (index, !stream.resource<*>{%c55296000}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%14}
%16 = arith.muli %0, %c34560 : index
%17 = arith.muli %16, %1 : index
%18 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%9, %0, %1](%15[%c0 to %14 for %14], %9, %0, %1) : (!stream.resource<*>{%14}, index, index, index) -> !stream.resource<*>{%17}
%19 = stream.async.transfer %18 : !stream.resource<*>{%17} -> !stream.resource<external>{%17}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%17} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
}
// -----// IR Dump After IPO (iree-util-ipo) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%5 = stream.async.transfer %4 : !stream.resource<external>{%3} -> !stream.resource<*>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%7 = stream.async.transfer %6 : !stream.resource<external>{%c55296000} -> !stream.resource<*>{%c55296000}
%8 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %c55296000 for %c55296000]) : (!stream.resource<*>{%c55296000}) -> !stream.resource<*>{%c55296000}
%9 = affine.apply #map()[%1]
%10 = arith.muli %0, %c102400 : index
%11 = arith.muli %10, %9 : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %9](%5[%c0 to %3 for %3], %1, %0, %9) : (!stream.resource<*>{%3}, index, index, index) -> !stream.resource<*>{%11}
%13 = arith.muli %0, %c276480 : index
%14 = arith.muli %13, %9 : index
%15 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %9](%1, %8[%c0 to %c55296000 for %c55296000], %12[%c0 to %11 for %11], %0, %9) : (index, !stream.resource<*>{%c55296000}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%14}
%16 = arith.muli %0, %c34560 : index
%17 = arith.muli %16, %1 : index
%18 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%9, %0, %1](%15[%c0 to %14 for %14], %9, %0, %1) : (!stream.resource<*>{%14}, index, index, index) -> !stream.resource<*>{%17}
%19 = stream.async.transfer %18 : !stream.resource<*>{%17} -> !stream.resource<external>{%17}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%17} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
}
// -----// IR Dump After VerifyLoweringToAsyncResourcesPass (iree-stream-verify-lowering-to-async-resources) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%5 = stream.async.transfer %4 : !stream.resource<external>{%3} -> !stream.resource<*>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%7 = stream.async.transfer %6 : !stream.resource<external>{%c55296000} -> !stream.resource<*>{%c55296000}
%8 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %c55296000 for %c55296000]) : (!stream.resource<*>{%c55296000}) -> !stream.resource<*>{%c55296000}
%9 = affine.apply #map()[%1]
%10 = arith.muli %0, %c102400 : index
%11 = arith.muli %10, %9 : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %9](%5[%c0 to %3 for %3], %1, %0, %9) : (!stream.resource<*>{%3}, index, index, index) -> !stream.resource<*>{%11}
%13 = arith.muli %0, %c276480 : index
%14 = arith.muli %13, %9 : index
%15 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %9](%1, %8[%c0 to %c55296000 for %c55296000], %12[%c0 to %11 for %11], %0, %9) : (index, !stream.resource<*>{%c55296000}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%14}
%16 = arith.muli %0, %c34560 : index
%17 = arith.muli %16, %1 : index
%18 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%9, %0, %1](%15[%c0 to %14 for %14], %9, %0, %1) : (!stream.resource<*>{%14}, index, index, index) -> !stream.resource<*>{%17}
%19 = stream.async.transfer %18 : !stream.resource<*>{%17} -> !stream.resource<external>{%17}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%17} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
}
// -----// IR Dump After MaterializeCopyOnWritePass (iree-stream-materialize-copy-on-write) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%5 = stream.async.transfer %4 : !stream.resource<external>{%3} -> !stream.resource<*>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%7 = stream.async.transfer %6 : !stream.resource<external>{%c55296000} -> !stream.resource<*>{%c55296000}
%8 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %c55296000 for %c55296000]) : (!stream.resource<*>{%c55296000}) -> !stream.resource<*>{%c55296000}
%9 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%10 = arith.muli %0, %c102400 : index
%11 = arith.muli %10, %9 : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %9](%5[%c0 to %3 for %3], %1, %0, %9) : (!stream.resource<*>{%3}, index, index, index) -> !stream.resource<*>{%11}
%13 = arith.muli %0, %c276480 : index
%14 = arith.muli %13, %9 : index
%15 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %9](%1, %8[%c0 to %c55296000 for %c55296000], %12[%c0 to %11 for %11], %0, %9) : (index, !stream.resource<*>{%c55296000}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%14}
%16 = arith.muli %0, %c34560 : index
%17 = arith.muli %16, %1 : index
%18 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%9, %0, %1](%15[%c0 to %14 for %14], %9, %0, %1) : (!stream.resource<*>{%14}, index, index, index) -> !stream.resource<*>{%17}
%19 = stream.async.transfer %18 : !stream.resource<*>{%17} -> !stream.resource<external>{%17}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%17} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%5 = stream.async.transfer %4 : !stream.resource<external>{%3} -> !stream.resource<*>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%7 = stream.async.transfer %6 : !stream.resource<external>{%c55296000} -> !stream.resource<*>{%c55296000}
%8 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %c55296000 for %c55296000]) : (!stream.resource<*>{%c55296000}) -> !stream.resource<*>{%c55296000}
%9 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%10 = arith.muli %0, %c102400 : index
%11 = arith.muli %10, %9 : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %9](%5[%c0 to %3 for %3], %1, %0, %9) : (!stream.resource<*>{%3}, index, index, index) -> !stream.resource<*>{%11}
%13 = arith.muli %0, %c276480 : index
%14 = arith.muli %13, %9 : index
%15 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %9](%1, %8[%c0 to %c55296000 for %c55296000], %12[%c0 to %11 for %11], %0, %9) : (index, !stream.resource<*>{%c55296000}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%14}
%16 = arith.muli %0, %c34560 : index
%17 = arith.muli %16, %1 : index
%18 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%9, %0, %1](%15[%c0 to %14 for %14], %9, %0, %1) : (!stream.resource<*>{%14}, index, index, index) -> !stream.resource<*>{%17}
%19 = stream.async.transfer %18 : !stream.resource<*>{%17} -> !stream.resource<external>{%17}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%17} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
// -----// IR Dump After ElideAsyncCopiesPass (iree-stream-elide-async-copies) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%5 = stream.async.transfer %4 : !stream.resource<external>{%3} -> !stream.resource<*>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%7 = stream.async.transfer %6 : !stream.resource<external>{%c55296000} -> !stream.resource<*>{%c55296000}
%8 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %c55296000 for %c55296000]) : (!stream.resource<*>{%c55296000}) -> !stream.resource<*>{%c55296000}
%9 = affine.apply #map()[%1]
%10 = arith.muli %0, %c102400 : index
%11 = arith.muli %10, %9 : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %9](%5[%c0 to %3 for %3], %1, %0, %9) : (!stream.resource<*>{%3}, index, index, index) -> !stream.resource<*>{%11}
%13 = arith.muli %0, %c276480 : index
%14 = arith.muli %13, %9 : index
%15 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %9](%1, %8[%c0 to %c55296000 for %c55296000], %12[%c0 to %11 for %11], %0, %9) : (index, !stream.resource<*>{%c55296000}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%14}
%16 = arith.muli %0, %c34560 : index
%17 = arith.muli %16, %1 : index
%18 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%9, %0, %1](%15[%c0 to %14 for %14], %9, %0, %1) : (!stream.resource<*>{%14}, index, index, index) -> !stream.resource<*>{%17}
%19 = stream.async.transfer %18 : !stream.resource<*>{%17} -> !stream.resource<external>{%17}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%17} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%5 = stream.async.transfer %4 : !stream.resource<external>{%3} -> !stream.resource<*>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%7 = stream.async.transfer %6 : !stream.resource<external>{%c55296000} -> !stream.resource<*>{%c55296000}
%8 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %c55296000 for %c55296000]) : (!stream.resource<*>{%c55296000}) -> !stream.resource<*>{%c55296000}
%9 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%10 = arith.muli %0, %c102400 : index
%11 = arith.muli %10, %9 : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %9](%5[%c0 to %3 for %3], %1, %0, %9) : (!stream.resource<*>{%3}, index, index, index) -> !stream.resource<*>{%11}
%13 = arith.muli %0, %c276480 : index
%14 = arith.muli %13, %9 : index
%15 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %9](%1, %8[%c0 to %c55296000 for %c55296000], %12[%c0 to %11 for %11], %0, %9) : (index, !stream.resource<*>{%c55296000}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%14}
%16 = arith.muli %0, %c34560 : index
%17 = arith.muli %16, %1 : index
%18 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%9, %0, %1](%15[%c0 to %14 for %14], %9, %0, %1) : (!stream.resource<*>{%14}, index, index, index) -> !stream.resource<*>{%17}
%19 = stream.async.transfer %18 : !stream.resource<*>{%17} -> !stream.resource<external>{%17}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%17} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
// -----// IR Dump After EmplaceAllocationsPass (iree-stream-emplace-allocations) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%5 = stream.async.transfer %4 : !stream.resource<external>{%3} -> !stream.resource<*>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%6 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%7 = stream.async.transfer %6 : !stream.resource<external>{%c55296000} -> !stream.resource<*>{%c55296000}
%8 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%7[%c0 to %c55296000 for %c55296000]) : (!stream.resource<*>{%c55296000}) -> !stream.resource<*>{%c55296000}
%9 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%10 = arith.muli %0, %c102400 : index
%11 = arith.muli %10, %9 : index
%12 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %9](%5[%c0 to %3 for %3], %1, %0, %9) : (!stream.resource<*>{%3}, index, index, index) -> !stream.resource<*>{%11}
%13 = arith.muli %0, %c276480 : index
%14 = arith.muli %13, %9 : index
%15 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %9](%1, %8[%c0 to %c55296000 for %c55296000], %12[%c0 to %11 for %11], %0, %9) : (index, !stream.resource<*>{%c55296000}, !stream.resource<*>{%11}, index, index) -> !stream.resource<*>{%14}
%16 = arith.muli %0, %c34560 : index
%17 = arith.muli %16, %1 : index
%18 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%9, %0, %1](%15[%c0 to %14 for %14], %9, %0, %1) : (!stream.resource<*>{%14}, index, index, index) -> !stream.resource<*>{%17}
%19 = stream.async.transfer %18 : !stream.resource<*>{%17} -> !stream.resource<external>{%17}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%17} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
// -----// IR Dump After RefineUsagePass (iree-stream-refine-usage) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%5[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%7 = affine.apply #map()[%1]
%8 = arith.muli %0, %c102400 : index
%9 = arith.muli %8, %7 : index
%10 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %7](%4[%c0 to %3 for %3], %1, %0, %7) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%9}
%11 = arith.muli %0, %c276480 : index
%12 = arith.muli %11, %7 : index
%13 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %7](%1, %6[%c0 to %c55296000 for %c55296000], %10[%c0 to %9 for %9], %0, %7) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%9}, index, index) -> !stream.resource<transient>{%12}
%14 = arith.muli %0, %c34560 : index
%15 = arith.muli %14, %1 : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%7, %0, %1](%13[%c0 to %12 for %12], %7, %0, %1) : (!stream.resource<transient>{%12}, index, index, index) -> !stream.resource<external>{%15}
%17 = stream.tensor.export %16 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %17 : !hal.buffer_view
}
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%5[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = arith.muli %0, %c102400 : index
%9 = arith.muli %8, %7 : index
%10 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %7](%4[%c0 to %3 for %3], %1, %0, %7) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%9}
%11 = arith.muli %0, %c276480 : index
%12 = arith.muli %11, %7 : index
%13 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %7](%1, %6[%c0 to %c55296000 for %c55296000], %10[%c0 to %9 for %9], %0, %7) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%9}, index, index) -> !stream.resource<transient>{%12}
%14 = arith.muli %0, %c34560 : index
%15 = arith.muli %14, %1 : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%7, %0, %1](%13[%c0 to %12 for %12], %7, %0, %1) : (!stream.resource<transient>{%12}, index, index, index) -> !stream.resource<external>{%15}
%17 = stream.tensor.export %16 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %17 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%5[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = arith.muli %0, %c102400 : index
%9 = arith.muli %8, %7 : index
%10 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %7](%4[%c0 to %3 for %3], %1, %0, %7) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%9}
%11 = arith.muli %0, %c276480 : index
%12 = arith.muli %11, %7 : index
%13 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %7](%1, %6[%c0 to %c55296000 for %c55296000], %10[%c0 to %9 for %9], %0, %7) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%9}, index, index) -> !stream.resource<transient>{%12}
%14 = arith.muli %0, %c34560 : index
%15 = arith.muli %14, %1 : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%7, %0, %1](%13[%c0 to %12 for %12], %7, %0, %1) : (!stream.resource<transient>{%12}, index, index, index) -> !stream.resource<external>{%15}
%17 = stream.tensor.export %16 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %17 : !hal.buffer_view
}
// -----// IR Dump After SimplifyGlobalAccesses (iree-util-simplify-global-accesses) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%5[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%7 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%8 = arith.muli %0, %c102400 : index
%9 = arith.muli %8, %7 : index
%10 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %7](%4[%c0 to %3 for %3], %1, %0, %7) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%9}
%11 = arith.muli %0, %c276480 : index
%12 = arith.muli %11, %7 : index
%13 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %7](%1, %6[%c0 to %c55296000 for %c55296000], %10[%c0 to %9 for %9], %0, %7) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%9}, index, index) -> !stream.resource<transient>{%12}
%14 = arith.muli %0, %c34560 : index
%15 = arith.muli %14, %1 : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%7, %0, %1](%13[%c0 to %12 for %12], %7, %0, %1) : (!stream.resource<transient>{%12}, index, index, index) -> !stream.resource<external>{%15}
%17 = stream.tensor.export %16 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %17 : !hal.buffer_view
}
// -----// IR Dump After ApplyPatterns (iree-util-apply-patterns) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%5[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%7 = affine.apply #map()[%1]
%8 = arith.muli %0, %c102400 : index
%9 = arith.muli %8, %7 : index
%10 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %7](%4[%c0 to %3 for %3], %1, %0, %7) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%9}
%11 = arith.muli %0, %c276480 : index
%12 = arith.muli %11, %7 : index
%13 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %7](%1, %6[%c0 to %c55296000 for %c55296000], %10[%c0 to %9 for %9], %0, %7) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%9}, index, index) -> !stream.resource<transient>{%12}
%14 = arith.muli %0, %c34560 : index
%15 = arith.muli %14, %1 : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%7, %0, %1](%13[%c0 to %12 for %12], %7, %0, %1) : (!stream.resource<transient>{%12}, index, index, index) -> !stream.resource<external>{%15}
%17 = stream.tensor.export %16 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %17 : !hal.buffer_view
}
}
// -----// IR Dump After FoldGlobals (iree-util-fold-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%5[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%7 = affine.apply #map()[%1]
%8 = arith.muli %0, %c102400 : index
%9 = arith.muli %8, %7 : index
%10 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %7](%4[%c0 to %3 for %3], %1, %0, %7) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%9}
%11 = arith.muli %0, %c276480 : index
%12 = arith.muli %11, %7 : index
%13 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %7](%1, %6[%c0 to %c55296000 for %c55296000], %10[%c0 to %9 for %9], %0, %7) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%9}, index, index) -> !stream.resource<transient>{%12}
%14 = arith.muli %0, %c34560 : index
%15 = arith.muli %14, %1 : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%7, %0, %1](%13[%c0 to %12 for %12], %7, %0, %1) : (!stream.resource<transient>{%12}, index, index, index) -> !stream.resource<external>{%15}
%17 = stream.tensor.export %16 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %17 : !hal.buffer_view
}
}
// -----// IR Dump After FuseGlobals (iree-util-fuse-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%5[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%7 = affine.apply #map()[%1]
%8 = arith.muli %0, %c102400 : index
%9 = arith.muli %8, %7 : index
%10 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %7](%4[%c0 to %3 for %3], %1, %0, %7) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%9}
%11 = arith.muli %0, %c276480 : index
%12 = arith.muli %11, %7 : index
%13 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %7](%1, %6[%c0 to %c55296000 for %c55296000], %10[%c0 to %9 for %9], %0, %7) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%9}, index, index) -> !stream.resource<transient>{%12}
%14 = arith.muli %0, %c34560 : index
%15 = arith.muli %14, %1 : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%7, %0, %1](%13[%c0 to %12 for %12], %7, %0, %1) : (!stream.resource<transient>{%12}, index, index, index) -> !stream.resource<external>{%15}
%17 = stream.tensor.export %16 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %17 : !hal.buffer_view
}
}
// -----// IR Dump After IPO (iree-util-ipo) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%5[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%7 = affine.apply #map()[%1]
%8 = arith.muli %0, %c102400 : index
%9 = arith.muli %8, %7 : index
%10 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %7](%4[%c0 to %3 for %3], %1, %0, %7) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%9}
%11 = arith.muli %0, %c276480 : index
%12 = arith.muli %11, %7 : index
%13 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %7](%1, %6[%c0 to %c55296000 for %c55296000], %10[%c0 to %9 for %9], %0, %7) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%9}, index, index) -> !stream.resource<transient>{%12}
%14 = arith.muli %0, %c34560 : index
%15 = arith.muli %14, %1 : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%7, %0, %1](%13[%c0 to %12 for %12], %7, %0, %1) : (!stream.resource<transient>{%12}, index, index, index) -> !stream.resource<external>{%15}
%17 = stream.tensor.export %16 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %17 : !hal.buffer_view
}
}
// -----// IR Dump After VerifyAsyncAccessRangesPass (iree-stream-verify-async-access-ranges) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%5[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%7 = affine.apply #map()[%1]
%8 = arith.muli %0, %c102400 : index
%9 = arith.muli %8, %7 : index
%10 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %7](%4[%c0 to %3 for %3], %1, %0, %7) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%9}
%11 = arith.muli %0, %c276480 : index
%12 = arith.muli %11, %7 : index
%13 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %7](%1, %6[%c0 to %c55296000 for %c55296000], %10[%c0 to %9 for %9], %0, %7) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%9}, index, index) -> !stream.resource<transient>{%12}
%14 = arith.muli %0, %c34560 : index
%15 = arith.muli %14, %1 : index
%16 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%7, %0, %1](%13[%c0 to %12 for %12], %7, %0, %1) : (!stream.resource<transient>{%12}, index, index, index) -> !stream.resource<external>{%15}
%17 = stream.tensor.export %16 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%15} -> !hal.buffer_view
util.return %17 : !hal.buffer_view
}
}
// -----// IR Dump After ScheduleExecutionPass (iree-stream-schedule-execution) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%results, %result_timepoint = stream.async.execute with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}) -> !stream.resource<external>{%12} {
%15 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%arg2[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%16 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%arg3[%c0 to %3 for %3], %1, %0, %6) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%8}
%17 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %15[%c0 to %c55296000 for %c55296000], %16[%c0 to %8 for %8], %0, %6) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}, index, index) -> !stream.resource<transient>{%10}
%18 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%17[%c0 to %10 for %10], %6, %0, %1) : (!stream.resource<transient>{%10}, index, index, index) -> !stream.resource<external>{%12}
stream.yield %18 : !stream.resource<external>{%12}
} => !stream.timepoint
%13 = stream.timepoint.await %result_timepoint => %results : !stream.resource<external>{%12}
%14 = stream.tensor.export %13 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %14 : !hal.buffer_view
}
// -----// IR Dump After ScheduleConcurrencyPass (iree-stream-schedule-concurrency) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%results, %result_timepoint = stream.async.execute with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}) -> !stream.resource<external>{%12} {
%15:2 = stream.async.concurrent with(%arg2 as %arg4: !stream.resource<external>{%c55296000}, %arg3 as %arg5: !stream.resource<external>{%3}) -> (!stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}) {
%18 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%arg4[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%19 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%arg5[%c0 to %3 for %3], %1, %0, %6) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%8}
stream.yield %18, %19 : !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}
}
%16 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %15#0[%c0 to %c55296000 for %c55296000], %15#1[%c0 to %8 for %8], %0, %6) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}, index, index) -> !stream.resource<transient>{%10}
%17 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%16[%c0 to %10 for %10], %6, %0, %1) : (!stream.resource<transient>{%10}, index, index, index) -> !stream.resource<external>{%12}
stream.yield %17 : !stream.resource<external>{%12}
} => !stream.timepoint
%13 = stream.timepoint.await %result_timepoint => %results : !stream.resource<external>{%12}
%14 = stream.tensor.export %13 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %14 : !hal.buffer_view
}
// -----// IR Dump After PropagateTimepointsPass (iree-stream-propagate-timepoints) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply #map()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%13 = stream.timepoint.immediate => !stream.timepoint
%14 = stream.timepoint.immediate => !stream.timepoint
%15 = stream.timepoint.join max(%13, %14) => !stream.timepoint
%results, %result_timepoint = stream.async.execute await(%15) => with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}) -> !stream.resource<external>{%12} {
%18:2 = stream.async.concurrent with(%arg2 as %arg4: !stream.resource<external>{%c55296000}, %arg3 as %arg5: !stream.resource<external>{%3}) -> (!stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}) {
%21 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%arg4[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%22 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%arg5[%c0 to %3 for %3], %1, %0, %6) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%8}
stream.yield %21, %22 : !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}
}
%19 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %18#0[%c0 to %c55296000 for %c55296000], %18#1[%c0 to %8 for %8], %0, %6) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}, index, index) -> !stream.resource<transient>{%10}
%20 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%19[%c0 to %10 for %10], %6, %0, %1) : (!stream.resource<transient>{%10}, index, index, index) -> !stream.resource<external>{%12}
stream.yield %20 : !stream.resource<external>{%12}
} => !stream.timepoint
%16 = stream.timepoint.await %result_timepoint => %results : !stream.resource<external>{%12}
%17 = stream.tensor.export %16 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %17 : !hal.buffer_view
}
}
// -----// IR Dump After MaterializeBuiltinsPass (iree-stream-materialize-builtins) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply #map()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%13 = stream.timepoint.immediate => !stream.timepoint
%14 = stream.timepoint.immediate => !stream.timepoint
%15 = stream.timepoint.join max(%13, %14) => !stream.timepoint
%results, %result_timepoint = stream.async.execute await(%15) => with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}) -> !stream.resource<external>{%12} {
%18:2 = stream.async.concurrent with(%arg2 as %arg4: !stream.resource<external>{%c55296000}, %arg3 as %arg5: !stream.resource<external>{%3}) -> (!stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}) {
%21 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%arg4[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%22 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%arg5[%c0 to %3 for %3], %1, %0, %6) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%8}
stream.yield %21, %22 : !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}
}
%19 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %18#0[%c0 to %c55296000 for %c55296000], %18#1[%c0 to %8 for %8], %0, %6) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}, index, index) -> !stream.resource<transient>{%10}
%20 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%19[%c0 to %10 for %10], %6, %0, %1) : (!stream.resource<transient>{%10}, index, index, index) -> !stream.resource<external>{%12}
stream.yield %20 : !stream.resource<external>{%12}
} => !stream.timepoint
%16 = stream.timepoint.await %result_timepoint => %results : !stream.resource<external>{%12}
%17 = stream.tensor.export %16 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %17 : !hal.buffer_view
}
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%results, %result_timepoint = stream.async.execute with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}) -> !stream.resource<external>{%12} {
%15:2 = stream.async.concurrent with(%arg2 as %arg4: !stream.resource<external>{%c55296000}, %arg3 as %arg5: !stream.resource<external>{%3}) -> (!stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}) {
%18 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%arg4[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%19 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%arg5[%c0 to %3 for %3], %1, %0, %6) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%8}
stream.yield %18, %19 : !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}
}
%16 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %15#0[%c0 to %c55296000 for %c55296000], %15#1[%c0 to %8 for %8], %0, %6) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}, index, index) -> !stream.resource<transient>{%10}
%17 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%16[%c0 to %10 for %10], %6, %0, %1) : (!stream.resource<transient>{%10}, index, index, index) -> !stream.resource<external>{%12}
stream.yield %17 : !stream.resource<external>{%12}
} => !stream.timepoint
%13 = stream.timepoint.await %result_timepoint => %results : !stream.resource<external>{%12}
%14 = stream.tensor.export %13 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %14 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%results, %result_timepoint = stream.async.execute with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}) -> !stream.resource<external>{%12} {
%15:2 = stream.async.concurrent with(%arg2 as %arg4: !stream.resource<external>{%c55296000}, %arg3 as %arg5: !stream.resource<external>{%3}) -> (!stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}) {
%18 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%arg4[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%19 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%arg5[%c0 to %3 for %3], %1, %0, %6) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%8}
stream.yield %18, %19 : !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}
}
%16 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %15#0[%c0 to %c55296000 for %c55296000], %15#1[%c0 to %8 for %8], %0, %6) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}, index, index) -> !stream.resource<transient>{%10}
%17 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%16[%c0 to %10 for %10], %6, %0, %1) : (!stream.resource<transient>{%10}, index, index, index) -> !stream.resource<external>{%12}
stream.yield %17 : !stream.resource<external>{%12}
} => !stream.timepoint
%13 = stream.timepoint.await %result_timepoint => %results : !stream.resource<external>{%12}
%14 = stream.tensor.export %13 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %14 : !hal.buffer_view
}
// -----// IR Dump After SimplifyGlobalAccesses (iree-util-simplify-global-accesses) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%results, %result_timepoint = stream.async.execute with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}) -> !stream.resource<external>{%12} {
%15:2 = stream.async.concurrent with(%arg2 as %arg4: !stream.resource<external>{%c55296000}, %arg3 as %arg5: !stream.resource<external>{%3}) -> (!stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}) {
%18 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%arg4[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%19 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%arg5[%c0 to %3 for %3], %1, %0, %6) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%8}
stream.yield %18, %19 : !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}
}
%16 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %15#0[%c0 to %c55296000 for %c55296000], %15#1[%c0 to %8 for %8], %0, %6) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}, index, index) -> !stream.resource<transient>{%10}
%17 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%16[%c0 to %10 for %10], %6, %0, %1) : (!stream.resource<transient>{%10}, index, index, index) -> !stream.resource<external>{%12}
stream.yield %17 : !stream.resource<external>{%12}
} => !stream.timepoint
%13 = stream.timepoint.await %result_timepoint => %results : !stream.resource<external>{%12}
%14 = stream.tensor.export %13 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %14 : !hal.buffer_view
}
// -----// IR Dump After ApplyPatterns (iree-util-apply-patterns) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply #map()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%results, %result_timepoint = stream.async.execute with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}) -> !stream.resource<external>{%12} {
%15:2 = stream.async.concurrent with(%arg2 as %arg4: !stream.resource<external>{%c55296000}, %arg3 as %arg5: !stream.resource<external>{%3}) -> (!stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}) {
%18 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%arg4[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%19 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%arg5[%c0 to %3 for %3], %1, %0, %6) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%8}
stream.yield %18, %19 : !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}
}
%16 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %15#0[%c0 to %c55296000 for %c55296000], %15#1[%c0 to %8 for %8], %0, %6) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}, index, index) -> !stream.resource<transient>{%10}
%17 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%16[%c0 to %10 for %10], %6, %0, %1) : (!stream.resource<transient>{%10}, index, index, index) -> !stream.resource<external>{%12}
stream.yield %17 : !stream.resource<external>{%12}
} => !stream.timepoint
%13 = stream.timepoint.await %result_timepoint => %results : !stream.resource<external>{%12}
%14 = stream.tensor.export %13 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %14 : !hal.buffer_view
}
}
// -----// IR Dump After FoldGlobals (iree-util-fold-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply #map()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%results, %result_timepoint = stream.async.execute with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}) -> !stream.resource<external>{%12} {
%15:2 = stream.async.concurrent with(%arg2 as %arg4: !stream.resource<external>{%c55296000}, %arg3 as %arg5: !stream.resource<external>{%3}) -> (!stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}) {
%18 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%arg4[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%19 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%arg5[%c0 to %3 for %3], %1, %0, %6) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%8}
stream.yield %18, %19 : !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}
}
%16 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %15#0[%c0 to %c55296000 for %c55296000], %15#1[%c0 to %8 for %8], %0, %6) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}, index, index) -> !stream.resource<transient>{%10}
%17 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%16[%c0 to %10 for %10], %6, %0, %1) : (!stream.resource<transient>{%10}, index, index, index) -> !stream.resource<external>{%12}
stream.yield %17 : !stream.resource<external>{%12}
} => !stream.timepoint
%13 = stream.timepoint.await %result_timepoint => %results : !stream.resource<external>{%12}
%14 = stream.tensor.export %13 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %14 : !hal.buffer_view
}
}
// -----// IR Dump After FuseGlobals (iree-util-fuse-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply #map()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%results, %result_timepoint = stream.async.execute with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}) -> !stream.resource<external>{%12} {
%15:2 = stream.async.concurrent with(%arg2 as %arg4: !stream.resource<external>{%c55296000}, %arg3 as %arg5: !stream.resource<external>{%3}) -> (!stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}) {
%18 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%arg4[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%19 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%arg5[%c0 to %3 for %3], %1, %0, %6) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%8}
stream.yield %18, %19 : !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}
}
%16 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %15#0[%c0 to %c55296000 for %c55296000], %15#1[%c0 to %8 for %8], %0, %6) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}, index, index) -> !stream.resource<transient>{%10}
%17 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%16[%c0 to %10 for %10], %6, %0, %1) : (!stream.resource<transient>{%10}, index, index, index) -> !stream.resource<external>{%12}
stream.yield %17 : !stream.resource<external>{%12}
} => !stream.timepoint
%13 = stream.timepoint.await %result_timepoint => %results : !stream.resource<external>{%12}
%14 = stream.tensor.export %13 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %14 : !hal.buffer_view
}
}
// -----// IR Dump After IPO (iree-util-ipo) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply #map()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%results, %result_timepoint = stream.async.execute with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}) -> !stream.resource<external>{%12} {
%15:2 = stream.async.concurrent with(%arg2 as %arg4: !stream.resource<external>{%c55296000}, %arg3 as %arg5: !stream.resource<external>{%3}) -> (!stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}) {
%18 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%arg4[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%19 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%arg5[%c0 to %3 for %3], %1, %0, %6) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%8}
stream.yield %18, %19 : !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}
}
%16 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %15#0[%c0 to %c55296000 for %c55296000], %15#1[%c0 to %8 for %8], %0, %6) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}, index, index) -> !stream.resource<transient>{%10}
%17 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%16[%c0 to %10 for %10], %6, %0, %1) : (!stream.resource<transient>{%10}, index, index, index) -> !stream.resource<external>{%12}
stream.yield %17 : !stream.resource<external>{%12}
} => !stream.timepoint
%13 = stream.timepoint.await %result_timepoint => %results : !stream.resource<external>{%12}
%14 = stream.tensor.export %13 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %14 : !hal.buffer_view
}
}
// -----// IR Dump After VerifyLoweringToAsyncPass (iree-stream-verify-lowering-to-async) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply #map()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%results, %result_timepoint = stream.async.execute with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}) -> !stream.resource<external>{%12} {
%15:2 = stream.async.concurrent with(%arg2 as %arg4: !stream.resource<external>{%c55296000}, %arg3 as %arg5: !stream.resource<external>{%3}) -> (!stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}) {
%18 = stream.async.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16(%arg4[%c0 to %c55296000 for %c55296000]) : (!stream.resource<external>{%c55296000}) -> !stream.resource<transient>{%c55296000}
%19 = stream.async.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%arg5[%c0 to %3 for %3], %1, %0, %6) : (!stream.resource<external>{%3}, index, index, index) -> !stream.resource<transient>{%8}
stream.yield %18, %19 : !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}
}
%16 = stream.async.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %15#0[%c0 to %c55296000 for %c55296000], %15#1[%c0 to %8 for %8], %0, %6) : (index, !stream.resource<transient>{%c55296000}, !stream.resource<transient>{%8}, index, index) -> !stream.resource<transient>{%10}
%17 = stream.async.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%16[%c0 to %10 for %10], %6, %0, %1) : (!stream.resource<transient>{%10}, index, index, index) -> !stream.resource<external>{%12}
stream.yield %17 : !stream.resource<external>{%12}
} => !stream.timepoint
%13 = stream.timepoint.await %result_timepoint => %results : !stream.resource<external>{%12}
%14 = stream.tensor.export %13 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %14 : !hal.buffer_view
}
}
// -----// IR Dump After ScheduleAllocationPass (iree-stream-schedule-allocation) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply #map()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%c0_0 = arith.constant 0 : index
%result, %result_timepoint = stream.resource.alloca uninitialized : !stream.resource<external>{%12} => !stream.timepoint
%13:4 = stream.resource.pack slices({
[0, 1] = %c55296000,
[0, 1] = %8,
[1, 2] = %10
}) : index
%result_1, %result_timepoint_2 = stream.resource.alloca uninitialized : !stream.resource<transient>{%13#0} => !stream.timepoint
%14 = stream.timepoint.join max(%result_timepoint, %result_timepoint_2) => !stream.timepoint
%15 = stream.cmd.execute await(%14) => with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}, %result as %arg4: !stream.resource<external>{%12}, %result_1 as %arg5: !stream.resource<transient>{%13#0}) {
stream.cmd.concurrent {
stream.cmd.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16 {
ro %arg2[%c0 for %c55296000] : !stream.resource<external>{%c55296000},
wo %arg5[%13#1 for %c55296000] : !stream.resource<transient>{%13#0}
}
stream.cmd.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg3[%c0 for %3] : !stream.resource<external>{%3},
wo %arg5[%13#2 for %8] : !stream.resource<transient>{%13#0}
}
}
stream.cmd.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg5[%13#1 for %c55296000] : !stream.resource<transient>{%13#0},
ro %arg5[%13#2 for %8] : !stream.resource<transient>{%13#0},
wo %arg5[%13#3 for %10] : !stream.resource<transient>{%13#0}
}
stream.cmd.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%6, %0, %1 : index, index, index) {
ro %arg5[%13#3 for %10] : !stream.resource<transient>{%13#0},
wo %arg4[%c0_0 for %12] : !stream.resource<external>{%12}
}
} => !stream.timepoint
%16 = stream.resource.dealloca await(%15) => %result_1 : !stream.resource<transient>{%13#0} => !stream.timepoint
%17 = stream.timepoint.join max(%16, %15) => !stream.timepoint
%18 = stream.timepoint.await %17 => %result : !stream.resource<external>{%12}
%19 = stream.tensor.export %18 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %19 : !hal.buffer_view
}
}
// -----// IR Dump After PackConstantsPass (iree-stream-pack-constants) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%c0_0 = arith.constant 0 : index
%result, %result_timepoint = stream.resource.alloca uninitialized : !stream.resource<external>{%12} => !stream.timepoint
%13:4 = stream.resource.pack slices({
[0, 1] = %c55296000,
[0, 1] = %8,
[1, 2] = %10
}) : index
%result_1, %result_timepoint_2 = stream.resource.alloca uninitialized : !stream.resource<transient>{%13#0} => !stream.timepoint
%14 = stream.timepoint.join max(%result_timepoint, %result_timepoint_2) => !stream.timepoint
%15 = stream.cmd.execute await(%14) => with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}, %result as %arg4: !stream.resource<external>{%12}, %result_1 as %arg5: !stream.resource<transient>{%13#0}) {
stream.cmd.concurrent {
stream.cmd.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16 {
ro %arg2[%c0 for %c55296000] : !stream.resource<external>{%c55296000},
wo %arg5[%13#1 for %c55296000] : !stream.resource<transient>{%13#0}
}
stream.cmd.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg3[%c0 for %3] : !stream.resource<external>{%3},
wo %arg5[%13#2 for %8] : !stream.resource<transient>{%13#0}
}
}
stream.cmd.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg5[%13#1 for %c55296000] : !stream.resource<transient>{%13#0},
ro %arg5[%13#2 for %8] : !stream.resource<transient>{%13#0},
wo %arg5[%13#3 for %10] : !stream.resource<transient>{%13#0}
}
stream.cmd.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%6, %0, %1 : index, index, index) {
ro %arg5[%13#3 for %10] : !stream.resource<transient>{%13#0},
wo %arg4[%c0_0 for %12] : !stream.resource<external>{%12}
}
} => !stream.timepoint
%16 = stream.resource.dealloca await(%15) => %result_1 : !stream.resource<transient>{%13#0} => !stream.timepoint
%17 = stream.timepoint.join max(%16, %15) => !stream.timepoint
%18 = stream.timepoint.await %17 => %result : !stream.resource<external>{%12}
%19 = stream.tensor.export %18 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %19 : !hal.buffer_view
}
// -----// IR Dump After LayoutSlicesPass (iree-stream-layout-slices) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%c0_0 = arith.constant 0 : index
%result, %result_timepoint = stream.resource.alloca uninitialized : !stream.resource<external>{%12} => !stream.timepoint
%c0_1 = arith.constant 0 : index
%c55296000_2 = arith.constant 55296000 : index
%c55296000_3 = arith.constant 55296000 : index
%c64 = arith.constant 64 : index
%13 = arith.addi %8, %c55296000_3 : index
%c64_4 = arith.constant 64 : index
%c64_5 = arith.constant 64 : index
%14 = arith.addi %13, %10 : index
%c64_6 = arith.constant 64 : index
%c64_7 = arith.constant 64 : index
%result_8, %result_timepoint_9 = stream.resource.alloca uninitialized : !stream.resource<transient>{%14} => !stream.timepoint
%15 = stream.timepoint.join max(%result_timepoint, %result_timepoint_9) => !stream.timepoint
%16 = stream.cmd.execute await(%15) => with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}, %result as %arg4: !stream.resource<external>{%12}, %result_8 as %arg5: !stream.resource<transient>{%14}) {
stream.cmd.concurrent {
stream.cmd.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16 {
ro %arg2[%c0 for %c55296000] : !stream.resource<external>{%c55296000},
wo %arg5[%c0_1 for %c55296000] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg3[%c0 for %3] : !stream.resource<external>{%3},
wo %arg5[%c55296000_3 for %8] : !stream.resource<transient>{%14}
}
}
stream.cmd.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg5[%c0_1 for %c55296000] : !stream.resource<transient>{%14},
ro %arg5[%c55296000_3 for %8] : !stream.resource<transient>{%14},
wo %arg5[%13 for %10] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%6, %0, %1 : index, index, index) {
ro %arg5[%13 for %10] : !stream.resource<transient>{%14},
wo %arg4[%c0_0 for %12] : !stream.resource<external>{%12}
}
} => !stream.timepoint
%17 = stream.resource.dealloca await(%16) => %result_8 : !stream.resource<transient>{%14} => !stream.timepoint
%18 = stream.timepoint.join max(%17, %16) => !stream.timepoint
%19 = stream.timepoint.await %18 => %result : !stream.resource<external>{%12}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
// -----// IR Dump After PropagateSubranges (iree-util-propagate-subranges) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply #map()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%c0_0 = arith.constant 0 : index
%result, %result_timepoint = stream.resource.alloca uninitialized : !stream.resource<external>{%12} => !stream.timepoint
%c0_1 = arith.constant 0 : index
%c55296000_2 = arith.constant 55296000 : index
%c55296000_3 = arith.constant 55296000 : index
%c64 = arith.constant 64 : index
%13 = arith.addi %8, %c55296000_3 : index
%c64_4 = arith.constant 64 : index
%c64_5 = arith.constant 64 : index
%14 = arith.addi %13, %10 : index
%c64_6 = arith.constant 64 : index
%c64_7 = arith.constant 64 : index
%result_8, %result_timepoint_9 = stream.resource.alloca uninitialized : !stream.resource<transient>{%14} => !stream.timepoint
%15 = stream.timepoint.join max(%result_timepoint, %result_timepoint_9) => !stream.timepoint
%16 = stream.cmd.execute await(%15) => with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}, %result as %arg4: !stream.resource<external>{%12}, %result_8 as %arg5: !stream.resource<transient>{%14}) {
stream.cmd.concurrent {
stream.cmd.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16 {
ro %arg2[%c0 for %c55296000] : !stream.resource<external>{%c55296000},
wo %arg5[%c0_1 for %c55296000] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg3[%c0 for %3] : !stream.resource<external>{%3},
wo %arg5[%c55296000_3 for %8] : !stream.resource<transient>{%14}
}
}
stream.cmd.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg5[%c0_1 for %c55296000] : !stream.resource<transient>{%14},
ro %arg5[%c55296000_3 for %8] : !stream.resource<transient>{%14},
wo %arg5[%13 for %10] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%6, %0, %1 : index, index, index) {
ro %arg5[%13 for %10] : !stream.resource<transient>{%14},
wo %arg4[%c0_0 for %12] : !stream.resource<external>{%12}
}
} => !stream.timepoint
%17 = stream.resource.dealloca await(%16) => %result_8 : !stream.resource<transient>{%14} => !stream.timepoint
%18 = stream.timepoint.join max(%17, %16) => !stream.timepoint
%19 = stream.timepoint.await %18 => %result : !stream.resource<external>{%12}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%result, %result_timepoint = stream.resource.alloca uninitialized : !stream.resource<external>{%12} => !stream.timepoint
%13 = arith.addi %8, %c55296000 : index
%14 = arith.addi %13, %10 : index
%result_0, %result_timepoint_1 = stream.resource.alloca uninitialized : !stream.resource<transient>{%14} => !stream.timepoint
%15 = stream.timepoint.join max(%result_timepoint, %result_timepoint_1) => !stream.timepoint
%16 = stream.cmd.execute await(%15) => with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}, %result as %arg4: !stream.resource<external>{%12}, %result_0 as %arg5: !stream.resource<transient>{%14}) {
stream.cmd.concurrent {
stream.cmd.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16 {
ro %arg2[%c0 for %c55296000] : !stream.resource<external>{%c55296000},
wo %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg3[%c0 for %3] : !stream.resource<external>{%3},
wo %arg5[%c55296000 for %8] : !stream.resource<transient>{%14}
}
}
stream.cmd.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14},
ro %arg5[%c55296000 for %8] : !stream.resource<transient>{%14},
wo %arg5[%13 for %10] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%6, %0, %1 : index, index, index) {
ro %arg5[%13 for %10] : !stream.resource<transient>{%14},
wo %arg4[%c0 for %12] : !stream.resource<external>{%12}
}
} => !stream.timepoint
%17 = stream.resource.dealloca await(%16) => %result_0 : !stream.resource<transient>{%14} => !stream.timepoint
%18 = stream.timepoint.join max(%17, %16) => !stream.timepoint
%19 = stream.timepoint.await %18 => %result : !stream.resource<external>{%12}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%result, %result_timepoint = stream.resource.alloca uninitialized : !stream.resource<external>{%12} => !stream.timepoint
%13 = arith.addi %8, %c55296000 : index
%14 = arith.addi %13, %10 : index
%result_0, %result_timepoint_1 = stream.resource.alloca uninitialized : !stream.resource<transient>{%14} => !stream.timepoint
%15 = stream.timepoint.join max(%result_timepoint, %result_timepoint_1) => !stream.timepoint
%16 = stream.cmd.execute await(%15) => with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}, %result as %arg4: !stream.resource<external>{%12}, %result_0 as %arg5: !stream.resource<transient>{%14}) {
stream.cmd.concurrent {
stream.cmd.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16 {
ro %arg2[%c0 for %c55296000] : !stream.resource<external>{%c55296000},
wo %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg3[%c0 for %3] : !stream.resource<external>{%3},
wo %arg5[%c55296000 for %8] : !stream.resource<transient>{%14}
}
}
stream.cmd.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14},
ro %arg5[%c55296000 for %8] : !stream.resource<transient>{%14},
wo %arg5[%13 for %10] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%6, %0, %1 : index, index, index) {
ro %arg5[%13 for %10] : !stream.resource<transient>{%14},
wo %arg4[%c0 for %12] : !stream.resource<external>{%12}
}
} => !stream.timepoint
%17 = stream.resource.dealloca await(%16) => %result_0 : !stream.resource<transient>{%14} => !stream.timepoint
%18 = stream.timepoint.join max(%17, %16) => !stream.timepoint
%19 = stream.timepoint.await %18 => %result : !stream.resource<external>{%12}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
// -----// IR Dump After SimplifyGlobalAccesses (iree-util-simplify-global-accesses) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%result, %result_timepoint = stream.resource.alloca uninitialized : !stream.resource<external>{%12} => !stream.timepoint
%13 = arith.addi %8, %c55296000 : index
%14 = arith.addi %13, %10 : index
%result_0, %result_timepoint_1 = stream.resource.alloca uninitialized : !stream.resource<transient>{%14} => !stream.timepoint
%15 = stream.timepoint.join max(%result_timepoint, %result_timepoint_1) => !stream.timepoint
%16 = stream.cmd.execute await(%15) => with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}, %result as %arg4: !stream.resource<external>{%12}, %result_0 as %arg5: !stream.resource<transient>{%14}) {
stream.cmd.concurrent {
stream.cmd.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16 {
ro %arg2[%c0 for %c55296000] : !stream.resource<external>{%c55296000},
wo %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg3[%c0 for %3] : !stream.resource<external>{%3},
wo %arg5[%c55296000 for %8] : !stream.resource<transient>{%14}
}
}
stream.cmd.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14},
ro %arg5[%c55296000 for %8] : !stream.resource<transient>{%14},
wo %arg5[%13 for %10] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%6, %0, %1 : index, index, index) {
ro %arg5[%13 for %10] : !stream.resource<transient>{%14},
wo %arg4[%c0 for %12] : !stream.resource<external>{%12}
}
} => !stream.timepoint
%17 = stream.resource.dealloca await(%16) => %result_0 : !stream.resource<transient>{%14} => !stream.timepoint
%18 = stream.timepoint.join max(%17, %16) => !stream.timepoint
%19 = stream.timepoint.await %18 => %result : !stream.resource<external>{%12}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
// -----// IR Dump After ApplyPatterns (iree-util-apply-patterns) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply #map()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%result, %result_timepoint = stream.resource.alloca uninitialized : !stream.resource<external>{%12} => !stream.timepoint
%13 = arith.addi %8, %c55296000 : index
%14 = arith.addi %13, %10 : index
%result_0, %result_timepoint_1 = stream.resource.alloca uninitialized : !stream.resource<transient>{%14} => !stream.timepoint
%15 = stream.timepoint.join max(%result_timepoint, %result_timepoint_1) => !stream.timepoint
%16 = stream.cmd.execute await(%15) => with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}, %result as %arg4: !stream.resource<external>{%12}, %result_0 as %arg5: !stream.resource<transient>{%14}) {
stream.cmd.concurrent {
stream.cmd.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16 {
ro %arg2[%c0 for %c55296000] : !stream.resource<external>{%c55296000},
wo %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg3[%c0 for %3] : !stream.resource<external>{%3},
wo %arg5[%c55296000 for %8] : !stream.resource<transient>{%14}
}
}
stream.cmd.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14},
ro %arg5[%c55296000 for %8] : !stream.resource<transient>{%14},
wo %arg5[%13 for %10] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%6, %0, %1 : index, index, index) {
ro %arg5[%13 for %10] : !stream.resource<transient>{%14},
wo %arg4[%c0 for %12] : !stream.resource<external>{%12}
}
} => !stream.timepoint
%17 = stream.resource.dealloca await(%16) => %result_0 : !stream.resource<transient>{%14} => !stream.timepoint
%18 = stream.timepoint.join max(%17, %16) => !stream.timepoint
%19 = stream.timepoint.await %18 => %result : !stream.resource<external>{%12}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
}
// -----// IR Dump After FoldGlobals (iree-util-fold-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply #map()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%result, %result_timepoint = stream.resource.alloca uninitialized : !stream.resource<external>{%12} => !stream.timepoint
%13 = arith.addi %8, %c55296000 : index
%14 = arith.addi %13, %10 : index
%result_0, %result_timepoint_1 = stream.resource.alloca uninitialized : !stream.resource<transient>{%14} => !stream.timepoint
%15 = stream.timepoint.join max(%result_timepoint, %result_timepoint_1) => !stream.timepoint
%16 = stream.cmd.execute await(%15) => with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}, %result as %arg4: !stream.resource<external>{%12}, %result_0 as %arg5: !stream.resource<transient>{%14}) {
stream.cmd.concurrent {
stream.cmd.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16 {
ro %arg2[%c0 for %c55296000] : !stream.resource<external>{%c55296000},
wo %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg3[%c0 for %3] : !stream.resource<external>{%3},
wo %arg5[%c55296000 for %8] : !stream.resource<transient>{%14}
}
}
stream.cmd.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14},
ro %arg5[%c55296000 for %8] : !stream.resource<transient>{%14},
wo %arg5[%13 for %10] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%6, %0, %1 : index, index, index) {
ro %arg5[%13 for %10] : !stream.resource<transient>{%14},
wo %arg4[%c0 for %12] : !stream.resource<external>{%12}
}
} => !stream.timepoint
%17 = stream.resource.dealloca await(%16) => %result_0 : !stream.resource<transient>{%14} => !stream.timepoint
%18 = stream.timepoint.join max(%17, %16) => !stream.timepoint
%19 = stream.timepoint.await %18 => %result : !stream.resource<external>{%12}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
}
// -----// IR Dump After FuseGlobals (iree-util-fuse-globals) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply #map()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%result, %result_timepoint = stream.resource.alloca uninitialized : !stream.resource<external>{%12} => !stream.timepoint
%13 = arith.addi %8, %c55296000 : index
%14 = arith.addi %13, %10 : index
%result_0, %result_timepoint_1 = stream.resource.alloca uninitialized : !stream.resource<transient>{%14} => !stream.timepoint
%15 = stream.timepoint.join max(%result_timepoint, %result_timepoint_1) => !stream.timepoint
%16 = stream.cmd.execute await(%15) => with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}, %result as %arg4: !stream.resource<external>{%12}, %result_0 as %arg5: !stream.resource<transient>{%14}) {
stream.cmd.concurrent {
stream.cmd.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16 {
ro %arg2[%c0 for %c55296000] : !stream.resource<external>{%c55296000},
wo %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg3[%c0 for %3] : !stream.resource<external>{%3},
wo %arg5[%c55296000 for %8] : !stream.resource<transient>{%14}
}
}
stream.cmd.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14},
ro %arg5[%c55296000 for %8] : !stream.resource<transient>{%14},
wo %arg5[%13 for %10] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%6, %0, %1 : index, index, index) {
ro %arg5[%13 for %10] : !stream.resource<transient>{%14},
wo %arg4[%c0 for %12] : !stream.resource<external>{%12}
}
} => !stream.timepoint
%17 = stream.resource.dealloca await(%16) => %result_0 : !stream.resource<transient>{%14} => !stream.timepoint
%18 = stream.timepoint.join max(%17, %16) => !stream.timepoint
%19 = stream.timepoint.await %18 => %result : !stream.resource<external>{%12}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
}
// -----// IR Dump After IPO (iree-util-ipo) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply #map()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%result, %result_timepoint = stream.resource.alloca uninitialized : !stream.resource<external>{%12} => !stream.timepoint
%13 = arith.addi %8, %c55296000 : index
%14 = arith.addi %13, %10 : index
%result_0, %result_timepoint_1 = stream.resource.alloca uninitialized : !stream.resource<transient>{%14} => !stream.timepoint
%15 = stream.timepoint.join max(%result_timepoint, %result_timepoint_1) => !stream.timepoint
%16 = stream.cmd.execute await(%15) => with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}, %result as %arg4: !stream.resource<external>{%12}, %result_0 as %arg5: !stream.resource<transient>{%14}) {
stream.cmd.concurrent {
stream.cmd.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16 {
ro %arg2[%c0 for %c55296000] : !stream.resource<external>{%c55296000},
wo %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg3[%c0 for %3] : !stream.resource<external>{%3},
wo %arg5[%c55296000 for %8] : !stream.resource<transient>{%14}
}
}
stream.cmd.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14},
ro %arg5[%c55296000 for %8] : !stream.resource<transient>{%14},
wo %arg5[%13 for %10] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%6, %0, %1 : index, index, index) {
ro %arg5[%13 for %10] : !stream.resource<transient>{%14},
wo %arg4[%c0 for %12] : !stream.resource<external>{%12}
}
} => !stream.timepoint
%17 = stream.resource.dealloca await(%16) => %result_0 : !stream.resource<transient>{%14} => !stream.timepoint
%18 = stream.timepoint.join max(%17, %16) => !stream.timepoint
%19 = stream.timepoint.await %18 => %result : !stream.resource<external>{%12}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
}
// -----// IR Dump After VerifyLoweringToCmdPass (iree-stream-verify-lowering-to-cmd) //----- //
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64", {cpu = "generic", cpu_features = "", 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", native_vector_size = 16 : i64, target_triple = "x86_64-unknown-unknown-eabi-elf"}>
#map = affine_map<()[s0] -> (s0 ceildiv 8)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#device_target_local = #hal.device.target<"local", [#executable_target_embedded_elf_x86_64_]>
module attributes {hal.device.targets = [#device_target_local]} {
stream.executable private @matmul_broad_dispatch_0 {
stream.executable.export public @matmul_broad_dispatch_0_pack_f16 workgroups() -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_0_pack_f16(%arg0: !stream.binding, %arg1: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>>
%1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
%2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [8640, 3200], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<8640x3200xf16>> -> tensor<8640x3200xf16>
%3 = tensor.empty() : tensor<2160x3200x4x1xf16>
%pack = tensor.pack %2 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [4, 1] into %3 : tensor<8640x3200xf16> -> tensor<2160x3200x4x1xf16>
flow.dispatch.tensor.store %pack, %1, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : tensor<2160x3200x4x1xf16> -> !flow.dispatch.tensor<writeonly:tensor<2160x3200x4x1xf16>>
return
}
}
}
stream.executable private @matmul_broad_dispatch_1 {
stream.executable.export public @matmul_broad_dispatch_1_pack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_1_pack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0], sizes = [%1, %0, 3200], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200xf32>>{%1, %0} -> tensor<?x?x3200xf32>
%6 = affine.apply #map()[%0]
%7 = tensor.empty(%1, %6) : tensor<?x?x3200x8x1xf32>
%pack = tensor.pack %5 padding_value(%cst : f32) outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %7 : tensor<?x?x3200xf32> -> tensor<?x?x3200x8x1xf32>
flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : tensor<?x?x3200x8x1xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_2 {
stream.executable.export public @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32(%arg0: index, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: index, %arg4: index, %arg5: !stream.binding) {
%cst = arith.constant 0.000000e+00 : f32
%c0 = arith.constant 0 : index
%0 = stream.binding.subspan %arg1[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>>
%1 = flow.dispatch.workload.ordinal %arg3, 1 : index
%2 = flow.dispatch.workload.ordinal %arg4, 2 : index
%3 = stream.binding.subspan %arg2[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2}
%4 = stream.binding.subspan %arg5[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
%5 = flow.dispatch.workload.ordinal %arg0, 0 : index
%6 = flow.dispatch.tensor.load %0, offsets = [0, 0, 0, 0], sizes = [2160, 3200, 4, 1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<2160x3200x4x1xf16>> -> tensor<2160x3200x4x1xf16>
%7 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 3200, 8, 1], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x3200x8x1xf32>>{%1, %2} -> tensor<?x?x3200x8x1xf32>
%8 = affine.apply #map()[%5]
%9 = tensor.empty(%1) : tensor<?x2160x3200x4x1xf16>
%10 = tensor.empty(%1, %8) : tensor<?x?x2160x8x4xf32>
%11 = linalg.fill ins(%cst : f32) outs(%10 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
%12 = linalg.generic {indexing_maps = [#map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%6 : tensor<2160x3200x4x1xf16>) outs(%9 : tensor<?x2160x3200x4x1xf16>) {
^bb0(%in: f16, %out: f16):
linalg.yield %in : f16
} -> tensor<?x2160x3200x4x1xf16>
%13 = linalg.batch_mmt4d ins(%7, %12 : tensor<?x?x3200x8x1xf32>, tensor<?x2160x3200x4x1xf16>) outs(%11 : tensor<?x?x2160x8x4xf32>) -> tensor<?x?x2160x8x4xf32>
flow.dispatch.tensor.store %13, %4, offsets = [0, 0, 0, 0, 0], sizes = [%1, %2, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : tensor<?x?x2160x8x4xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x2160x8x4xf32>>{%1, %2}
return
}
}
}
stream.executable private @matmul_broad_dispatch_3 {
stream.executable.export public @matmul_broad_dispatch_3_unpack_f32 workgroups(%arg0: index, %arg1: index, %arg2: index) -> (index, index, index) {
%x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg0, %arg1, %arg2
stream.return %x, %y, %z : index, index, index
}
builtin.module {
func.func @matmul_broad_dispatch_3_unpack_f32(%arg0: !stream.binding, %arg1: index, %arg2: index, %arg3: index, %arg4: !stream.binding) {
%c0 = arith.constant 0 : index
%0 = flow.dispatch.workload.ordinal %arg1, 0 : index
%1 = flow.dispatch.workload.ordinal %arg2, 1 : index
%2 = flow.dispatch.workload.ordinal %arg3, 2 : index
%3 = stream.binding.subspan %arg0[%c0] : !stream.binding -> !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0}
%4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
%5 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0, 0], sizes = [%1, %0, 2160, 8, 4], strides = [1, 1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x2160x8x4xf32>>{%1, %0} -> tensor<?x?x2160x8x4xf32>
%6 = tensor.empty(%1, %2) : tensor<?x?x8640xf32>
%unpack = tensor.unpack %5 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 4] into %6 : tensor<?x?x2160x8x4xf32> -> tensor<?x?x8640xf32>
flow.dispatch.tensor.store %unpack, %4, offsets = [0, 0, 0], sizes = [%1, %2, 8640], strides = [1, 1, 1] : tensor<?x?x8640xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x8640xf32>>{%1, %2}
return
}
}
}
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply #map()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%result, %result_timepoint = stream.resource.alloca uninitialized : !stream.resource<external>{%12} => !stream.timepoint
%13 = arith.addi %8, %c55296000 : index
%14 = arith.addi %13, %10 : index
%result_0, %result_timepoint_1 = stream.resource.alloca uninitialized : !stream.resource<transient>{%14} => !stream.timepoint
%15 = stream.timepoint.join max(%result_timepoint, %result_timepoint_1) => !stream.timepoint
%16 = stream.cmd.execute await(%15) => with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}, %result as %arg4: !stream.resource<external>{%12}, %result_0 as %arg5: !stream.resource<transient>{%14}) {
stream.cmd.concurrent {
stream.cmd.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16 {
ro %arg2[%c0 for %c55296000] : !stream.resource<external>{%c55296000},
wo %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg3[%c0 for %3] : !stream.resource<external>{%3},
wo %arg5[%c55296000 for %8] : !stream.resource<transient>{%14}
}
}
stream.cmd.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14},
ro %arg5[%c55296000 for %8] : !stream.resource<transient>{%14},
wo %arg5[%13 for %10] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%6, %0, %1 : index, index, index) {
ro %arg5[%13 for %10] : !stream.resource<transient>{%14},
wo %arg4[%c0 for %12] : !stream.resource<external>{%12}
}
} => !stream.timepoint
%17 = stream.resource.dealloca await(%16) => %result_0 : !stream.resource<transient>{%14} => !stream.timepoint
%18 = stream.timepoint.join max(%17, %16) => !stream.timepoint
%19 = stream.timepoint.await %18 => %result : !stream.resource<external>{%12}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
}
// -----// IR Dump After Canonicalizer (canonicalize) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%result, %result_timepoint = stream.resource.alloca uninitialized : !stream.resource<external>{%12} => !stream.timepoint
%13 = arith.addi %8, %c55296000 : index
%14 = arith.addi %13, %10 : index
%result_0, %result_timepoint_1 = stream.resource.alloca uninitialized : !stream.resource<transient>{%14} => !stream.timepoint
%15 = stream.timepoint.join max(%result_timepoint, %result_timepoint_1) => !stream.timepoint
%16 = stream.cmd.execute await(%15) => with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}, %result as %arg4: !stream.resource<external>{%12}, %result_0 as %arg5: !stream.resource<transient>{%14}) {
stream.cmd.concurrent {
stream.cmd.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16 {
ro %arg2[%c0 for %c55296000] : !stream.resource<external>{%c55296000},
wo %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg3[%c0 for %3] : !stream.resource<external>{%3},
wo %arg5[%c55296000 for %8] : !stream.resource<transient>{%14}
}
}
stream.cmd.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14},
ro %arg5[%c55296000 for %8] : !stream.resource<transient>{%14},
wo %arg5[%13 for %10] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%6, %0, %1 : index, index, index) {
ro %arg5[%13 for %10] : !stream.resource<transient>{%14},
wo %arg4[%c0 for %12] : !stream.resource<external>{%12}
}
} => !stream.timepoint
%17 = stream.resource.dealloca await(%16) => %result_0 : !stream.resource<transient>{%14} => !stream.timepoint
%18 = stream.timepoint.join max(%17, %16) => !stream.timepoint
%19 = stream.timepoint.await %18 => %result : !stream.resource<external>{%12}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
// -----// IR Dump After CSE (cse) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c276480 : index
%10 = arith.muli %9, %6 : index
%11 = arith.muli %0, %c34560 : index
%12 = arith.muli %11, %1 : index
%result, %result_timepoint = stream.resource.alloca uninitialized : !stream.resource<external>{%12} => !stream.timepoint
%13 = arith.addi %8, %c55296000 : index
%14 = arith.addi %13, %10 : index
%result_0, %result_timepoint_1 = stream.resource.alloca uninitialized : !stream.resource<transient>{%14} => !stream.timepoint
%15 = stream.timepoint.join max(%result_timepoint, %result_timepoint_1) => !stream.timepoint
%16 = stream.cmd.execute await(%15) => with(%5 as %arg2: !stream.resource<external>{%c55296000}, %4 as %arg3: !stream.resource<external>{%3}, %result as %arg4: !stream.resource<external>{%12}, %result_0 as %arg5: !stream.resource<transient>{%14}) {
stream.cmd.concurrent {
stream.cmd.dispatch @matmul_broad_dispatch_0::@matmul_broad_dispatch_0_pack_f16 {
ro %arg2[%c0 for %c55296000] : !stream.resource<external>{%c55296000},
wo %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_1::@matmul_broad_dispatch_1_pack_f32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg3[%c0 for %3] : !stream.resource<external>{%3},
wo %arg5[%c55296000 for %8] : !stream.resource<transient>{%14}
}
}
stream.cmd.dispatch @matmul_broad_dispatch_2::@matmul_broad_dispatch_2_batch_mmt4d_DxDx2160x3200x8x4x1_f32xf16xf32[%1, %0, %6](%1, %0, %6 : index, index, index) {
ro %arg5[%c0 for %c55296000] : !stream.resource<transient>{%14},
ro %arg5[%c55296000 for %8] : !stream.resource<transient>{%14},
wo %arg5[%13 for %10] : !stream.resource<transient>{%14}
}
stream.cmd.dispatch @matmul_broad_dispatch_3::@matmul_broad_dispatch_3_unpack_f32[%6, %0, %1](%6, %0, %1 : index, index, index) {
ro %arg5[%13 for %10] : !stream.resource<transient>{%14},
wo %arg4[%c0 for %12] : !stream.resource<external>{%12}
}
} => !stream.timepoint
%17 = stream.resource.dealloca await(%16) => %result_0 : !stream.resource<transient>{%14} => !stream.timepoint
%18 = stream.timepoint.join max(%17, %16) => !stream.timepoint
%19 = stream.timepoint.await %18 => %result : !stream.resource<external>{%12}
%20 = stream.tensor.export %19 : tensor<?x?x8640xf32>{%0, %1} in !stream.resource<external>{%12} -> !hal.buffer_view
util.return %20 : !hal.buffer_view
}
// -----// IR Dump After SimplifyGlobalAccesses (iree-util-simplify-global-accesses) //----- //
util.func public @matmul_broad(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view) -> !hal.buffer_view attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "sync func @matmul_broad(%input0: tensor<?x?x3200xf32>, %input1: tensor<8640x3200xf16>) -> (%output0: tensor<?x?x8640xf32>)"}} {
%c34560 = arith.constant 34560 : index
%c276480 = arith.constant 276480 : index
%c102400 = arith.constant 102400 : index
%c55296000 = arith.constant 55296000 : index
%c12800 = arith.constant 12800 : index
%c0 = arith.constant 0 : index
%c8640 = arith.constant 8640 : index
%c3200 = arith.constant 3200 : index
%0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index
%1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index
%element_type_f32 = hal.element_type<f32> : i32
%dense_row_major = hal.encoding_type<dense_row_major> : i32
hal.buffer_view.assert<%arg0 : !hal.buffer_view> message("input0") shape([%0, %1, %c3200]) type(%element_type_f32) encoding(%dense_row_major)
%2 = arith.muli %0, %c12800 : index
%3 = arith.muli %2, %1 : index
%4 = stream.tensor.import %arg0 : !hal.buffer_view -> tensor<?x?x3200xf32>{%0, %1} in !stream.resource<external>{%3}
%element_type_f16 = hal.element_type<f16> : i32
hal.buffer_view.assert<%arg1 : !hal.buffer_view> message("input1") shape([%c8640, %c3200]) type(%element_type_f16) encoding(%dense_row_major)
%5 = stream.tensor.import %arg1 : !hal.buffer_view -> tensor<8640x3200xf16> in !stream.resource<external>{%c55296000}
%6 = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%1]
%7 = arith.muli %0, %c102400 : index
%8 = arith.muli %7, %6 : index
%9 = arith.muli %0, %c27648
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment