derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
2002 lines
92 KiB
MLIR
2002 lines
92 KiB
MLIR
// RUN: mlir-opt %s -transform-interpreter -cse -split-input-file | FileCheck %s
|
|
|
|
!vecA = vector<1x16x16x4xi8>
|
|
!vecB = vector<1x16x16x4xi8>
|
|
!vecC = vector<16x16xi32>
|
|
!memrefA = memref<1x32x16x4xi8>
|
|
!memrefB = memref<1x16x32x4xi8>
|
|
!memrefC = memref<32x32xi32>
|
|
#map = affine_map<(d0, d4, d1, d2, d3) -> (d0, d1, d3, d4)>
|
|
#map1 = affine_map<(d0, d4, d1, d2, d3) -> (d0, d3, d2, d4)>
|
|
#map2 = affine_map<(d0, d4, d1, d2, d3) -> (d1, d2)>
|
|
func.func @brgemm_int8(
|
|
%arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
|
|
{
|
|
%c0 = arith.constant 0 : index
|
|
%0 = ub.poison : i8
|
|
%32 = ub.poison : i32
|
|
|
|
%1 = vector.transfer_read %arg0[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
|
|
!memrefA, !vecA
|
|
%2 = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
|
|
!memrefB, !vecB
|
|
|
|
%3 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
|
|
%4 = vector.contract {
|
|
indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["reduction", "reduction", "parallel", "parallel", "reduction"],
|
|
kind = #vector.kind<add>}
|
|
%1, %2, %3 : !vecA, !vecB into !vecC
|
|
|
|
vector.transfer_write %4, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
|
|
return %arg2 : !memrefC
|
|
}
|
|
|
|
// CHECK-LABEL: @brgemm_int8
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x64xi8>
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x64xi8>
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x16xi32>
|
|
// CHECK: x86.amx.tile_muli
|
|
// CHECK: x86.amx.tile_store {{.*}} !x86.amx.tile<16x16xi32>
|
|
// CHECK-NOT: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<16x16x4xi8>
|
|
!vecB = vector<16x16x4xi8>
|
|
!vecC = vector<16x16xi32>
|
|
!memrefA = memref<32x16x4xi8>
|
|
!memrefB = memref<16x32x4xi8>
|
|
!memrefC = memref<32x32xi32>
|
|
#map = affine_map<(d4, d1, d2, d3) -> (d1, d3, d4)>
|
|
#map1 = affine_map<(d4, d1, d2, d3) -> (d3, d2, d4)>
|
|
#map2 = affine_map<(d4, d1, d2, d3) -> (d1, d2)>
|
|
func.func @matmul_int8(
|
|
%arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
|
|
{
|
|
%c0 = arith.constant 0 : index
|
|
%0 = ub.poison : i8
|
|
%32 = ub.poison : i32
|
|
|
|
%1 = vector.transfer_read %arg0[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} :
|
|
!memrefA, !vecA
|
|
%2 = vector.transfer_read %arg1[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} :
|
|
!memrefB, !vecB
|
|
|
|
%3 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
|
|
%4 = vector.contract {
|
|
indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["reduction", "parallel", "parallel", "reduction"],
|
|
kind = #vector.kind<add>}
|
|
%1, %2, %3 : !vecA, !vecB into !vecC
|
|
|
|
vector.transfer_write %4, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
|
|
return %arg2 : !memrefC
|
|
}
|
|
|
|
// CHECK-LABEL: @matmul_int8
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x64xi8>
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x64xi8>
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x16xi32>
|
|
// CHECK: x86.amx.tile_muli
|
|
// CHECK: x86.amx.tile_store {{.*}} !x86.amx.tile<16x16xi32>
|
|
// CHECK-NOT: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<1x16x16x2xbf16>
|
|
!vecB = vector<1x16x16x2xbf16>
|
|
!vecC = vector<16x16xf32>
|
|
!memrefA = memref<1x32x16x2xbf16>
|
|
!memrefB = memref<1x16x32x2xbf16>
|
|
!memrefC = memref<32x32xf32>
|
|
#map = affine_map<(d0, d4, d1, d2, d3) -> (d0, d1, d3, d4)>
|
|
#map1 = affine_map<(d0, d4, d1, d2, d3) -> (d0, d3, d2, d4)>
|
|
#map2 = affine_map<(d0, d4, d1, d2, d3) -> (d1, d2)>
|
|
func.func @brgemm_bf16(
|
|
%arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
|
|
{
|
|
%c0 = arith.constant 0 : index
|
|
%0 = ub.poison : bf16
|
|
%32 = ub.poison : f32
|
|
|
|
%1 = vector.transfer_read %arg0[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
|
|
!memrefA, !vecA
|
|
%2 = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
|
|
!memrefB, !vecB
|
|
|
|
%3 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
|
|
%4 = vector.contract {
|
|
indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["reduction", "reduction", "parallel", "parallel", "reduction"],
|
|
kind = #vector.kind<add>}
|
|
%1, %2, %3 : !vecA, !vecB into !vecC
|
|
|
|
vector.transfer_write %4, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
|
|
return %arg2 : !memrefC
|
|
}
|
|
|
|
// CHECK-LABEL: @brgemm_bf16
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x32xbf16>
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x32xbf16>
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x16xf32>
|
|
// CHECK: x86.amx.tile_mulf
|
|
// CHECK: x86.amx.tile_store {{.*}} !x86.amx.tile<16x16xf32>
|
|
// CHECK-NOT: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<1x16x16x2xbf16>
|
|
!vecB = vector<1x16x16x2xbf16>
|
|
!vecC = vector<1x16x16xf32>
|
|
!memrefA = memref<1x32x16x2xbf16>
|
|
!memrefB = memref<1x16x32x2xbf16>
|
|
!memrefC = memref<1x32x32xf32>
|
|
#map = affine_map<(d0, d4, d1, d2, d3) -> (d0, d1, d3, d4)>
|
|
#map1 = affine_map<(d0, d4, d1, d2, d3) -> (d0, d3, d2, d4)>
|
|
#map2 = affine_map<(d0, d4, d1, d2, d3) -> (d0, d1, d2)>
|
|
func.func @batch_matmul_bf16(
|
|
%arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
|
|
{
|
|
%c0 = arith.constant 0 : index
|
|
%0 = ub.poison : bf16
|
|
%32 = ub.poison : f32
|
|
|
|
%1 = vector.transfer_read %arg0[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
|
|
!memrefA, !vecA
|
|
%2 = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
|
|
!memrefB, !vecB
|
|
|
|
%3 = vector.transfer_read %arg2[%c0, %c0, %c0], %32 {in_bounds = [true, true, true]} : !memrefC, !vecC
|
|
|
|
%4 = vector.contract {
|
|
indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["parallel", "reduction", "parallel", "parallel", "reduction"],
|
|
kind = #vector.kind<add>}
|
|
%1, %2, %3 : !vecA, !vecB into !vecC
|
|
|
|
vector.transfer_write %4, %arg2[%c0, %c0, %c0] {in_bounds = [true, true, true]} : !vecC, !memrefC
|
|
|
|
return %arg2 : !memrefC
|
|
}
|
|
|
|
// CHECK-LABEL: @batch_matmul_bf16
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x32xbf16>
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x32xbf16>
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x16xf32>
|
|
// CHECK: x86.amx.tile_mulf
|
|
// CHECK: x86.amx.tile_store {{.*}} !x86.amx.tile<16x16xf32>
|
|
// CHECK-NOT: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<1x16x16x4xf8E5M2>
|
|
!vecB = vector<1x16x16x4xf8E5M2>
|
|
!vecC = vector<16x16xf32>
|
|
!memrefA = memref<1x32x16x4xf8E5M2>
|
|
!memrefB = memref<1x16x32x4xf8E5M2>
|
|
!memrefC = memref<32x32xf32>
|
|
#map = affine_map<(d0, d4, d1, d2, d3) -> (d0, d1, d3, d4)>
|
|
#map1 = affine_map<(d0, d4, d1, d2, d3) -> (d0, d3, d2, d4)>
|
|
#map2 = affine_map<(d0, d4, d1, d2, d3) -> (d1, d2)>
|
|
func.func @brgemm_f8E5M2(
|
|
%arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
|
|
{
|
|
%c0 = arith.constant 0 : index
|
|
%0 = ub.poison : f8E5M2
|
|
%32 = ub.poison : f32
|
|
|
|
%1 = vector.transfer_read %arg0[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
|
|
!memrefA, !vecA
|
|
%2 = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
|
|
!memrefB, !vecB
|
|
|
|
%3 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
|
|
%4 = vector.contract {
|
|
indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["reduction", "reduction", "parallel", "parallel", "reduction"],
|
|
kind = #vector.kind<add>}
|
|
%1, %2, %3 : !vecA, !vecB into !vecC
|
|
|
|
vector.transfer_write %4, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
|
|
return %arg2 : !memrefC
|
|
}
|
|
|
|
// CHECK-LABEL: @brgemm_f8E5M2
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x64xf8E5M2>
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x64xf8E5M2>
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x16xf32>
|
|
// CHECK: x86.amx.tile_mulf
|
|
// CHECK: x86.amx.tile_store {{.*}} !x86.amx.tile<16x16xf32>
|
|
// CHECK-NOT: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<16x64xi8>
|
|
!vecB = vector<64x16xi8>
|
|
!vecC = vector<16x16xi32>
|
|
!memrefA = memref<32x64xi8>
|
|
!memrefB = memref<64x32xi8>
|
|
!memrefC = memref<32x32xi32>
|
|
#map = affine_map<(d1, d2, d3) -> (d1, d3)>
|
|
#map1 = affine_map<(d1, d2, d3) -> (d3, d2)>
|
|
#map2 = affine_map<(d1, d2, d3) -> (d1, d2)>
|
|
func.func @online_packing_int8(
|
|
%arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
|
|
{
|
|
%c0 = arith.constant 0 : index
|
|
%0 = ub.poison : i8
|
|
%32 = ub.poison : i32
|
|
|
|
%1 = vector.transfer_read %arg0[%c0, %c0], %0 {in_bounds = [true, true]} :
|
|
!memrefA, !vecA
|
|
%2 = vector.transfer_read %arg1[%c0, %c0], %0 {in_bounds = [true, true]} :
|
|
!memrefB, !vecB
|
|
|
|
%3 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
|
|
%bias = arith.constant dense<13> : !vecC
|
|
|
|
%4 = vector.contract {
|
|
indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["parallel", "parallel", "reduction"],
|
|
kind = #vector.kind<add>}
|
|
%1, %2, %3 : !vecA, !vecB into !vecC
|
|
|
|
%5 = arith.addi %4, %bias : !vecC
|
|
|
|
vector.transfer_write %5, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
|
|
return %arg2 : !memrefC
|
|
}
|
|
|
|
// CHECK-LABEL: @online_packing_int8
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x64xi8>
|
|
// CHECK: scf.for
|
|
// CHECK: vector.shuffle{{.*}}[0, 16, 32, 48, 1, 17, 33, 49, 2, 18, 34, 50, 3, 19, 35, 51, 4, 20, 36, 52, 5, 21, 37, 53, 6, 22, 38, 54, 7, 23, 39, 55] : vector<32xi8>, vector<32xi8>
|
|
// CHECK-NEXT: vector.shuffle{{.*}}[8, 24, 40, 56, 9, 25, 41, 57, 10, 26, 42, 58, 11, 27, 43, 59, 12, 28, 44, 60, 13, 29, 45, 61, 14, 30, 46, 62, 15, 31, 47, 63] : vector<32xi8>, vector<32xi8>
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x64xi8>
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x16xi32>
|
|
// CHECK: x86.amx.tile_muli
|
|
// CHECK: x86.amx.tile_store {{.*}} !x86.amx.tile<16x16xi32>
|
|
// CHECK: vector.transfer_read
|
|
// CHECK: arith.addi
|
|
// CHECK: vector.transfer_write
|
|
// CHECK-NOT: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<1x16x32xbf16>
|
|
!vecB = vector<1x32x16xbf16>
|
|
!vecC = vector<16x16xf32>
|
|
!memrefA = memref<1x32x32xbf16>
|
|
!memrefB = memref<1x32x32xbf16>
|
|
!memrefC = memref<32x32xf32>
|
|
#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
|
|
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
|
|
#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
|
|
func.func @online_packing_bf16(
|
|
%arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
|
|
{
|
|
%c0 = arith.constant 0 : index
|
|
%0 = ub.poison : bf16
|
|
%32 = ub.poison : f32
|
|
|
|
%1 = vector.transfer_read %arg0[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} :
|
|
!memrefA, !vecA
|
|
%2 = vector.transfer_read %arg1[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} :
|
|
!memrefB, !vecB
|
|
|
|
%3 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
|
|
%4 = vector.contract {
|
|
indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["reduction", "parallel", "parallel", "reduction"],
|
|
kind = #vector.kind<add>}
|
|
%1, %2, %3 : !vecA, !vecB into !vecC
|
|
|
|
vector.transfer_write %4, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
|
|
return %arg2 : !memrefC
|
|
}
|
|
|
|
// CHECK-LABEL: @online_packing_bf16
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x32xbf16>
|
|
// CHECK: scf.for
|
|
// CHECK: vector.shuffle{{.*}}[0, 16, 1, 17, 2, 18, 3, 19, 4, 20, 5, 21, 6, 22, 7, 23] : vector<16xbf16>, vector<16xbf16>
|
|
// CHECK-NEXT: vector.shuffle{{.*}}[8, 24, 9, 25, 10, 26, 11, 27, 12, 28, 13, 29, 14, 30, 15, 31] : vector<16xbf16>, vector<16xbf16>
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x32xbf16>
|
|
// CHECK: x86.amx.tile_load {{.*}} !x86.amx.tile<16x16xf32>
|
|
// CHECK: x86.amx.tile_mulf
|
|
// CHECK: x86.amx.tile_store {{.*}} !x86.amx.tile<16x16xf32>
|
|
// CHECK-NOT: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecAB = vector<1x16x16x2xbf16>
|
|
!vecC = vector<16x16xf32>
|
|
!memrefA = memref<1x32x16x2xbf16, strided<[8192, 128, 2, 1], offset: ?>>
|
|
!memrefB = memref<1x16x32x2xbf16, strided<[16384, 256, 2, 1], offset: ?>>
|
|
!memrefC = memref<32x32xf32, strided<[128, 1], offset: ?>>
|
|
|
|
#map = affine_map<(d0, d1, d2, d3, d4) -> (d0, d2, d4, d1)>
|
|
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d4, d3, d1)>
|
|
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d2, d3)>
|
|
|
|
func.func @brgemm_bf16_loop(%arg0: memref<16x64x64x2xbf16>, %arg1: memref<16x64x128x2xbf16>, %arg2: memref<64x128xf32>) -> memref<64x128xf32> {
|
|
%0 = ub.poison : f32
|
|
%1 = ub.poison : bf16
|
|
%c0 = arith.constant 0 : index
|
|
%c64 = arith.constant 64 : index
|
|
%c128 = arith.constant 128 : index
|
|
%c16 = arith.constant 16 : index
|
|
%c32 = arith.constant 32 : index
|
|
%c1 = arith.constant 1 : index
|
|
|
|
scf.for %arg3 = %c0 to %c64 step %c32 {
|
|
scf.for %arg4 = %c0 to %c128 step %c32 {
|
|
|
|
%subview = memref.subview %arg2[%arg3, %arg4] [32, 32] [1, 1] :
|
|
memref<64x128xf32> to !memrefC
|
|
%2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
%3 = vector.transfer_read %subview[%c0, %c16], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
%4 = vector.transfer_read %subview[%c16, %c0], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
%5 = vector.transfer_read %subview[%c16, %c16], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
|
|
%6:4 = scf.for %arg5 = %c0 to %c16 step %c1 iter_args(%arg6 = %2, %arg7 = %3, %arg8 = %4, %arg9 = %5) -> (!vecC, !vecC, !vecC, !vecC) {
|
|
%7:4 = scf.for %arg10 = %c0 to %c64 step %c16 iter_args(%arg11 = %arg6, %arg12 = %arg7, %arg13 = %arg8, %arg14 = %arg9) -> (!vecC, !vecC, !vecC, !vecC) {
|
|
|
|
%subview_0 = memref.subview %arg0[%arg5, %arg3, %arg10, 0] [1, 32, 16, 2] [1, 1, 1, 1] :
|
|
memref<16x64x64x2xbf16> to !memrefA
|
|
%subview_1 = memref.subview %arg1[%arg5, %arg10, %arg4, 0] [1, 16, 32, 2] [1, 1, 1, 1] :
|
|
memref<16x64x128x2xbf16> to !memrefB
|
|
%8 = vector.transfer_read %subview_0[%c0, %c0, %c0, %c0], %1 {in_bounds = [true, true, true, true]} :
|
|
!memrefA, !vecAB
|
|
%9 = vector.transfer_read %subview_1[%c0, %c0, %c0, %c0], %1 {in_bounds = [true, true, true, true]} :
|
|
!memrefB, !vecAB
|
|
|
|
%10 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%8, %9, %arg11 {unroll_shape = array<i64: 1, 2, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
|
|
%11 = vector.transfer_read %subview_1[%c0, %c0, %c16, %c0], %1 {in_bounds = [true, true, true, true]} :
|
|
!memrefB, !vecAB
|
|
%12 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%8, %11, %arg12 {unroll_shape = array<i64: 1, 2, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
%13 = vector.transfer_read %subview_0[%c0, %c16, %c0, %c0], %1 {in_bounds = [true, true, true, true]} :
|
|
!memrefA, !vecAB
|
|
%14 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%13, %9, %arg13 {unroll_shape = array<i64: 1, 2, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
%15 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%13, %11, %arg14 {unroll_shape = array<i64: 1, 2, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
|
|
scf.yield %10, %12, %14, %15 : !vecC, !vecC, !vecC, !vecC
|
|
}
|
|
scf.yield %7#0, %7#1, %7#2, %7#3 : !vecC, !vecC, !vecC, !vecC
|
|
}
|
|
|
|
vector.transfer_write %6#3, %subview[%c16, %c16] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
vector.transfer_write %6#2, %subview[%c16, %c0] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
vector.transfer_write %6#1, %subview[%c0, %c16] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
vector.transfer_write %6#0, %subview[%c0, %c0] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
}
|
|
}
|
|
|
|
return %arg2 : memref<64x128xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: @brgemm_bf16_loop
|
|
// CHECK-2: scf.for {{.*}} -> (!x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>) {
|
|
// CHECK-4: x86.amx.tile_zero : !x86.amx.tile<16x16xf32>
|
|
// CHECK-4: x86.amx.tile_load
|
|
// CHECK-4: x86.amx.tile_mulf
|
|
// CHECK: scf.yield {{.*}} : !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>
|
|
// CHECK-NOT: scf.for {{.*}} vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>
|
|
// CHECK-NOT: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecAB = vector<16x16x4xi8>
|
|
!vecC = vector<16x16xi32>
|
|
!memrefA = memref<16x16x4xi8, strided<[256, 4, 1], offset: ?>>
|
|
!memrefB = memref<16x32x4xi8, strided<[512, 4, 1], offset: ?>>
|
|
!memrefC = memref<16x32xi32, strided<[128, 1], offset: ?>>
|
|
|
|
#map = affine_map<(d0, d1, d2, d3) -> (d1, d3, d0)>
|
|
#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d2, d0)>
|
|
#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
|
|
|
|
func.func @matmul_int8_loop(%arg0: memref<64x64x4xi8>, %arg1: memref<64x128x4xi8>, %arg2: memref<64x128xi32>) {
|
|
%0 = ub.poison : i32
|
|
%1 = ub.poison : i8
|
|
%c0 = arith.constant 0 : index
|
|
%c64 = arith.constant 64 : index
|
|
%c128 = arith.constant 128 : index
|
|
%c16 = arith.constant 16 : index
|
|
%c32 = arith.constant 32 : index
|
|
scf.for %arg3 = %c0 to %c64 step %c16 {
|
|
scf.for %arg4 = %c0 to %c128 step %c32 {
|
|
|
|
%subview = memref.subview %arg2[%arg3, %arg4] [16, 32] [1, 1] :
|
|
memref<64x128xi32> to !memrefC
|
|
%2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
%3 = vector.transfer_read %subview[%c0, %c16], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
|
|
%4:2 = scf.for %arg5 = %c0 to %c64 step %c16 iter_args(%arg6 = %2, %arg7 = %3) -> (!vecC, !vecC) {
|
|
|
|
%subview_0 = memref.subview %arg0[%arg3, %arg5, 0] [16, 16, 4] [1, 1, 1] :
|
|
memref<64x64x4xi8> to !memrefA
|
|
%subview_1 = memref.subview %arg1[%arg5, %arg4, 0] [16, 32, 4] [1, 1, 1] :
|
|
memref<64x128x4xi8> to !memrefB
|
|
%5 = vector.transfer_read %subview_0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
|
|
!memrefA, !vecAB
|
|
%6 = vector.transfer_read %subview_1[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
|
|
!memrefB, !vecAB
|
|
%7 = vector.transfer_read %subview_1[%c0, %c16, %c0], %1 {in_bounds = [true, true, true]} :
|
|
!memrefB, !vecAB
|
|
|
|
%8 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%5, %6, %arg6 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
%9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%5, %7, %arg7 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
|
|
scf.yield %8, %9 : !vecC, !vecC
|
|
}
|
|
|
|
vector.transfer_write %4#1, %subview[%c0, %c16] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
vector.transfer_write %4#0, %subview[%c0, %c0] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
}
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: @matmul_int8_loop
|
|
// CHECK-2: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
|
|
// CHECK: scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>) {
|
|
// CHECK-3: x86.amx.tile_load
|
|
// CHECK-2: x86.amx.tile_muli
|
|
// CHECK: scf.yield {{.*}} !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>
|
|
// CHECK-NOT: scf.for {{.*}} vector<16x16xi32>, vector<16x16xi32>
|
|
// CHECK-NOT: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecAB = vector<16x16x4xf8E4M3FN>
|
|
!vecC = vector<16x16xf32>
|
|
!memrefA = memref<16x16x4xf8E4M3FN, strided<[256, 4, 1], offset: ?>>
|
|
!memrefB = memref<16x32x4xf8E4M3FN, strided<[512, 4, 1], offset: ?>>
|
|
!memrefC = memref<16x32xf32, strided<[128, 1], offset: ?>>
|
|
|
|
#map = affine_map<(d0, d1, d2, d3) -> (d1, d3, d0)>
|
|
#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d2, d0)>
|
|
#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
|
|
|
|
func.func @matmul_f8E4M3FN_loop(%arg0: memref<64x64x4xf8E4M3FN>, %arg1: memref<64x128x4xf8E4M3FN>, %arg2: memref<64x128xf32>) {
|
|
%0 = ub.poison : f32
|
|
%1 = ub.poison : f8E4M3FN
|
|
%c0 = arith.constant 0 : index
|
|
%c64 = arith.constant 64 : index
|
|
%c128 = arith.constant 128 : index
|
|
%c16 = arith.constant 16 : index
|
|
%c32 = arith.constant 32 : index
|
|
scf.for %arg3 = %c0 to %c64 step %c16 {
|
|
scf.for %arg4 = %c0 to %c128 step %c32 {
|
|
|
|
%subview = memref.subview %arg2[%arg3, %arg4] [16, 32] [1, 1] :
|
|
memref<64x128xf32> to !memrefC
|
|
%2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
%3 = vector.transfer_read %subview[%c0, %c16], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
|
|
%4:2 = scf.for %arg5 = %c0 to %c64 step %c16 iter_args(%arg6 = %2, %arg7 = %3) -> (!vecC, !vecC) {
|
|
|
|
%subview_0 = memref.subview %arg0[%arg3, %arg5, 0] [16, 16, 4] [1, 1, 1] :
|
|
memref<64x64x4xf8E4M3FN> to !memrefA
|
|
%subview_1 = memref.subview %arg1[%arg5, %arg4, 0] [16, 32, 4] [1, 1, 1] :
|
|
memref<64x128x4xf8E4M3FN> to !memrefB
|
|
%5 = vector.transfer_read %subview_0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
|
|
!memrefA, !vecAB
|
|
%6 = vector.transfer_read %subview_1[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
|
|
!memrefB, !vecAB
|
|
%7 = vector.transfer_read %subview_1[%c0, %c16, %c0], %1 {in_bounds = [true, true, true]} :
|
|
!memrefB, !vecAB
|
|
|
|
%8 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%5, %6, %arg6 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
%9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%5, %7, %arg7 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
|
|
scf.yield %8, %9 : !vecC, !vecC
|
|
}
|
|
|
|
vector.transfer_write %4#1, %subview[%c0, %c16] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
vector.transfer_write %4#0, %subview[%c0, %c0] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
}
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: @matmul_f8E4M3FN_loop
|
|
// CHECK-2: x86.amx.tile_zero : !x86.amx.tile<16x16xf32>
|
|
// CHECK: scf.for {{.*}} -> (!x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>) {
|
|
// CHECK-3: x86.amx.tile_load
|
|
// CHECK-2: x86.amx.tile_mulf
|
|
// CHECK: scf.yield {{.*}} !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>
|
|
// CHECK-NOT: scf.for {{.*}} vector<16x16xf32>, vector<16x16xf32>
|
|
// CHECK-NOT: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecAB = vector<1x16x16x4xi8>
|
|
!vecC = vector<1x16x16xi32>
|
|
!memrefA = memref<1x16x16x4xi8, strided<[16384, 256, 4, 1], offset: ?>>
|
|
!memrefB = memref<1x16x32x4xi8, strided<[32768, 512, 4, 1], offset: ?>>
|
|
!memrefC = memref<1x16x32xi32, strided<[8192, 128, 1], offset: ?>>
|
|
|
|
#map = affine_map<(d0, d1, d2, d3, d4) -> (d0, d2, d4, d1)>
|
|
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d4, d3, d1)>
|
|
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d2, d3)>
|
|
|
|
func.func @batch_matmul_int8_loop(%arg0: memref<16x64x64x4xi8>, %arg1: memref<16x64x128x4xi8>, %arg2: memref<16x64x128xi32>) {
|
|
%0 = ub.poison : i32
|
|
%1 = ub.poison : i8
|
|
%c0 = arith.constant 0 : index
|
|
%c64 = arith.constant 64 : index
|
|
%c128 = arith.constant 128 : index
|
|
%c16 = arith.constant 16 : index
|
|
%c32 = arith.constant 32 : index
|
|
%c1 = arith.constant 1 : index
|
|
scf.for %arg3 = %c0 to %c64 step %c16 {
|
|
scf.for %arg4 = %c0 to %c128 step %c32 {
|
|
scf.for %arg5 = %c0 to %c16 step %c1 {
|
|
|
|
%subview = memref.subview %arg2[%arg5, %arg3, %arg4] [1, 16, 32] [1, 1, 1] :
|
|
memref<16x64x128xi32> to !memrefC
|
|
%2 = vector.transfer_read %subview[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} :
|
|
!memrefC, !vecC
|
|
%3 = vector.transfer_read %subview[%c0, %c0, %c16], %0 {in_bounds = [true, true, true]} :
|
|
!memrefC, !vecC
|
|
%4:2 = scf.for %arg6 = %c0 to %c64 step %c16 iter_args(%arg7 = %2, %arg8 = %3) -> (!vecC, !vecC) {
|
|
|
|
%subview_0 = memref.subview %arg0[%arg5, %arg3, %arg6, 0] [1, 16, 16, 4] [1, 1, 1, 1] :
|
|
memref<16x64x64x4xi8> to !memrefA
|
|
%subview_1 = memref.subview %arg1[%arg5, %arg6, %arg4, 0] [1, 16, 32, 4] [1, 1, 1, 1] :
|
|
memref<16x64x128x4xi8> to !memrefB
|
|
%5 = vector.transfer_read %subview_0[%c0, %c0, %c0, %c0], %1 {in_bounds = [true, true, true, true]} :
|
|
!memrefA, !vecAB
|
|
%6 = vector.transfer_read %subview_1[%c0, %c0, %c0, %c0], %1 {in_bounds = [true, true, true, true]} :
|
|
!memrefB, !vecAB
|
|
%7 = vector.transfer_read %subview_1[%c0, %c0, %c16, %c0], %1 {in_bounds = [true, true, true, true]} :
|
|
!memrefB, !vecAB
|
|
%8 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%5, %6, %arg7 {unroll_shape = array<i64: 1, 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
%9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%5, %7, %arg8 {unroll_shape = array<i64: 1, 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
scf.yield %8, %9 : !vecC, !vecC
|
|
}
|
|
|
|
vector.transfer_write %4#1, %subview[%c0, %c0, %c16] {in_bounds = [true, true, true]} :
|
|
!vecC, !memrefC
|
|
vector.transfer_write %4#0, %subview[%c0, %c0, %c0] {in_bounds = [true, true, true]} :
|
|
!vecC, !memrefC
|
|
}
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: @batch_matmul_int8_loop
|
|
// CHECK-2: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
|
|
// CHECK: scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>) {
|
|
// CHECK-3: x86.amx.tile_load
|
|
// CHECK-2: x86.amx.tile_muli
|
|
// CHECK: scf.yield {{.*}} !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>
|
|
// CHECK-NOT: scf.for {{.*}} vector<16x16xi32>, vector<16x16xi32>
|
|
// CHECK-NOT: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<1x16x32xbf16>
|
|
!vecB = vector<1x32x16xbf16>
|
|
!vecC = vector<16x16xf32>
|
|
!memrefA = memref<1x32x32xbf16, strided<[6144, 96, 1], offset: ?>>
|
|
!memrefB = memref<1x32x32xbf16, strided<[12288, 128, 1], offset: ?>>
|
|
!memrefC = memref<32x32xf32, strided<[128, 1], offset: ?>>
|
|
|
|
#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
|
|
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
|
|
#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
|
|
|
|
func.func @online_packing_bf16_loop(%arg0: memref<16x64x96xbf16>, %arg1: memref<16x96x128xbf16>, %arg2: memref<64x128xf32>) -> memref<64x128xf32> {
|
|
%0 = ub.poison : f32
|
|
%1 = ub.poison : bf16
|
|
%c0 = arith.constant 0 : index
|
|
%c64 = arith.constant 64 : index
|
|
%c128 = arith.constant 128 : index
|
|
%c16 = arith.constant 16 : index
|
|
%c96 = arith.constant 96 : index
|
|
%c32 = arith.constant 32 : index
|
|
%c1 = arith.constant 1 : index
|
|
scf.for %arg3 = %c0 to %c64 step %c32 {
|
|
scf.for %arg4 = %c0 to %c128 step %c32 {
|
|
|
|
%subview = memref.subview %arg2[%arg3, %arg4] [32, 32] [1, 1] :
|
|
memref<64x128xf32> to !memrefC
|
|
%2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
%3 = vector.transfer_read %subview[%c0, %c16], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
%4 = vector.transfer_read %subview[%c16, %c0], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
%5 = vector.transfer_read %subview[%c16, %c16], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
|
|
%6:4 = scf.for %arg5 = %c0 to %c16 step %c1 iter_args(%arg6 = %2, %arg7 = %3, %arg8 = %4, %arg9 = %5) -> (!vecC, !vecC, !vecC, !vecC) {
|
|
%7:4 = scf.for %arg10 = %c0 to %c96 step %c32 iter_args(%arg11 = %arg6, %arg12 = %arg7, %arg13 = %arg8, %arg14 = %arg9) -> (!vecC, !vecC, !vecC, !vecC) {
|
|
|
|
%subview_0 = memref.subview %arg0[%arg5, %arg3, %arg10] [1, 32, 32] [1, 1, 1] :
|
|
memref<16x64x96xbf16> to !memrefA
|
|
%subview_1 = memref.subview %arg1[%arg5, %arg10, %arg4] [1, 32, 32] [1, 1, 1] :
|
|
memref<16x96x128xbf16> to !memrefB
|
|
%8 = vector.transfer_read %subview_0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
|
|
!memrefA, !vecA
|
|
%9 = vector.transfer_read %subview_0[%c0, %c16, %c0], %1 {in_bounds = [true, true, true]} :
|
|
!memrefA, !vecA
|
|
%10 = vector.transfer_read %subview_1[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
|
|
!memrefB, !vecB
|
|
%11 = vector.transfer_read %subview_1[%c0, %c0, %c16], %1 {in_bounds = [true, true, true]} :
|
|
!memrefB, !vecB
|
|
|
|
%12 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%8, %10, %arg11 {unroll_shape = array<i64: 1, 16, 16, 32>} : !vecA, !vecB into !vecC
|
|
%13 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%8, %11, %arg12 {unroll_shape = array<i64: 1, 16, 16, 32>} : !vecA, !vecB into !vecC
|
|
%14 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%9, %10, %arg13 {unroll_shape = array<i64: 1, 16, 16, 32>} : !vecA, !vecB into !vecC
|
|
%15 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%9, %11, %arg14 {unroll_shape = array<i64: 1, 16, 16, 32>} : !vecA, !vecB into !vecC
|
|
|
|
scf.yield %12, %13, %14, %15 : !vecC, !vecC, !vecC, !vecC
|
|
}
|
|
scf.yield %7#0, %7#1, %7#2, %7#3 : !vecC, !vecC, !vecC, !vecC
|
|
}
|
|
vector.transfer_write %6#3, %subview[%c16, %c16] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
vector.transfer_write %6#2, %subview[%c16, %c0] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
vector.transfer_write %6#1, %subview[%c0, %c16] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
vector.transfer_write %6#0, %subview[%c0, %c0] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
}
|
|
}
|
|
%alloc = memref.alloc() : memref<64x128xf32>
|
|
memref.copy %arg2, %alloc : memref<64x128xf32> to memref<64x128xf32>
|
|
return %alloc : memref<64x128xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: @online_packing_bf16_loop
|
|
// CHECK-COUNT-4: x86.amx.tile_zero : !x86.amx.tile<16x16xf32>
|
|
// CHECK-COUNT-4: scf.for {{.*}} -> (!x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>) {
|
|
// CHECK: vector.shuffle{{.*}}[0, 32, 1, 33, 2, 34, 3, 35, 8, 40, 9, 41, 10, 42, 11, 43, 16, 48, 17, 49, 18, 50, 19, 51, 24, 56, 25, 57, 26, 58, 27, 59] : vector<32xbf16>, vector<32xbf16>
|
|
// CHECK-NEXT: vector.shuffle{{.*}}[4, 36, 5, 37, 6, 38, 7, 39, 12, 44, 13, 45, 14, 46, 15, 47, 20, 52, 21, 53, 22, 54, 23, 55, 28, 60, 29, 61, 30, 62, 31, 63] : vector<32xbf16>, vector<32xbf16>
|
|
// CHECK: x86.amx.tile_load
|
|
// CHECK: x86.amx.tile_mulf
|
|
// CHECK: scf.yield {{.*}} : !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>
|
|
// CHECK: vector.shuffle{{.*}}[0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23] : vector<16xf32>, vector<16xf32>
|
|
// CHECK-NEXT: vector.shuffle{{.*}}[8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31] : vector<16xf32>, vector<16xf32>
|
|
// CHECK-NOT: scf.for {{.*}} vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>
|
|
// CHECK-NOT: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<16x32xbf16>
|
|
!vecB = vector<32x16xbf16>
|
|
!vecC = vector<16x16xf32>
|
|
!memrefA = memref<32x32xbf16, strided<[96, 1], offset: ?>>
|
|
!memrefB = memref<32x32xbf16, strided<[128, 1], offset: ?>>
|
|
!memrefC = memref<32x32xf32, strided<[128, 1], offset: ?>>
|
|
|
|
#map = affine_map<(d1, d2, d3) -> (d1, d3)>
|
|
#map1 = affine_map<(d1, d2, d3) -> (d3, d2)>
|
|
#map2 = affine_map<(d1, d2, d3) -> (d1, d2)>
|
|
|
|
func.func @online_packing_bf16_loop_lb_non_zero(%arg0: memref<64x96xbf16>, %arg1: memref<96x128xbf16>, %arg2: memref<64x128xf32>, %klb: index, %kub: index) -> memref<64x128xf32> {
|
|
%0 = ub.poison : f32
|
|
%1 = ub.poison : bf16
|
|
%c0 = arith.constant 0 : index
|
|
%c64 = arith.constant 64 : index
|
|
%c128 = arith.constant 128 : index
|
|
%c16 = arith.constant 16 : index
|
|
%c96 = arith.constant 96 : index
|
|
%c32 = arith.constant 32 : index
|
|
%c1 = arith.constant 1 : index
|
|
scf.for %arg3 = %c0 to %c64 step %c32 {
|
|
scf.for %arg4 = %c0 to %c128 step %c32 {
|
|
|
|
%subview = memref.subview %arg2[%arg3, %arg4] [32, 32] [1, 1] :
|
|
memref<64x128xf32> to !memrefC
|
|
%2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
%3 = vector.transfer_read %subview[%c0, %c16], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
%4 = vector.transfer_read %subview[%c16, %c0], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
%5 = vector.transfer_read %subview[%c16, %c16], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
|
|
%7:4 = scf.for %arg10 = %klb to %kub step %c32 iter_args(%arg11 = %2, %arg12 = %3, %arg13 = %4, %arg14 = %5) -> (!vecC, !vecC, !vecC, !vecC) {
|
|
%subview_0 = memref.subview %arg0[%arg3, %arg10] [32, 32] [1, 1] :
|
|
memref<64x96xbf16> to !memrefA
|
|
%subview_1 = memref.subview %arg1[%arg10, %arg4] [32, 32] [1, 1] :
|
|
memref<96x128xbf16> to !memrefB
|
|
%8 = vector.transfer_read %subview_0[%c0, %c0], %1 {in_bounds = [true, true]} :
|
|
!memrefA, !vecA
|
|
%9 = vector.transfer_read %subview_0[%c16, %c0], %1 {in_bounds = [true, true]} :
|
|
!memrefA, !vecA
|
|
%10 = vector.transfer_read %subview_1[%c0, %c0], %1 {in_bounds = [true, true]} :
|
|
!memrefB, !vecB
|
|
%11 = vector.transfer_read %subview_1[%c0, %c16], %1 {in_bounds = [true, true]} :
|
|
!memrefB, !vecB
|
|
|
|
%12 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%8, %10, %arg11 {unroll_shape = array<i64: 1, 16, 16, 32>} : !vecA, !vecB into !vecC
|
|
%13 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%8, %11, %arg12 {unroll_shape = array<i64: 1, 16, 16, 32>} : !vecA, !vecB into !vecC
|
|
%14 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%9, %10, %arg13 {unroll_shape = array<i64: 1, 16, 16, 32>} : !vecA, !vecB into !vecC
|
|
%15 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%9, %11, %arg14 {unroll_shape = array<i64: 1, 16, 16, 32>} : !vecA, !vecB into !vecC
|
|
|
|
scf.yield %12, %13, %14, %15 : !vecC, !vecC, !vecC, !vecC
|
|
}
|
|
vector.transfer_write %7#3, %subview[%c16, %c16] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
vector.transfer_write %7#2, %subview[%c16, %c0] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
vector.transfer_write %7#1, %subview[%c0, %c16] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
vector.transfer_write %7#0, %subview[%c0, %c0] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
}
|
|
}
|
|
%alloc = memref.alloc() : memref<64x128xf32>
|
|
memref.copy %arg2, %alloc : memref<64x128xf32> to memref<64x128xf32>
|
|
return %alloc : memref<64x128xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: @online_packing_bf16_loop_lb_non_zero
|
|
// CHECK-COUNT-4: x86.amx.tile_zero : !x86.amx.tile<16x16xf32>
|
|
// CHECK: scf.for {{.*}} -> (!x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>) {
|
|
// CHECK: vector.shuffle{{.*}}[0, 32, 1, 33, 2, 34, 3, 35, 8, 40, 9, 41, 10, 42, 11, 43, 16, 48, 17, 49, 18, 50, 19, 51, 24, 56, 25, 57, 26, 58, 27, 59] : vector<32xbf16>, vector<32xbf16>
|
|
// CHECK-NEXT: vector.shuffle{{.*}}[4, 36, 5, 37, 6, 38, 7, 39, 12, 44, 13, 45, 14, 46, 15, 47, 20, 52, 21, 53, 22, 54, 23, 55, 28, 60, 29, 61, 30, 62, 31, 63] : vector<32xbf16>, vector<32xbf16>
|
|
// CHECK: x86.amx.tile_load
|
|
// CHECK: x86.amx.tile_mulf
|
|
// CHECK: scf.yield {{.*}} : !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>
|
|
// CHECK: vector.shuffle{{.*}}[0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23] : vector<16xf32>, vector<16xf32>
|
|
// CHECK-NEXT: vector.shuffle{{.*}}[8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31] : vector<16xf32>, vector<16xf32>
|
|
// CHECK-NOT: scf.for {{.*}} vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>
|
|
// CHECK-NOT: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
|
|
// -----
|
|
|
|
!vecA = vector<16x64xf8E5M2>
|
|
!vecB = vector<64x16xf8E5M2>
|
|
!vecC = vector<16x16xf32>
|
|
!memrefA = memref<32x64xf8E5M2, strided<[256, 1], offset: ?>>
|
|
!memrefB = memref<64x32xf8E5M2, strided<[128, 1], offset: ?>>
|
|
!memrefC = memref<32x32xf32, strided<[128, 1], offset: ?>>
|
|
|
|
#map = affine_map<(d0, d1, d2) -> (d0, d2)>
|
|
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
|
|
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
func.func @online_packing_f8E5M2_matmul_loop(%arg0: memref<64x256xf8E5M2>, %arg1: memref<256x128xf8E5M2>, %arg2: memref<64x128xf32>) -> memref<64x128xf32> {
|
|
%c16 = arith.constant 16 : index
|
|
%0 = ub.poison : f32
|
|
%1 = ub.poison : f8E5M2
|
|
%c0 = arith.constant 0 : index
|
|
%c64 = arith.constant 64 : index
|
|
%c128 = arith.constant 128 : index
|
|
%c256 = arith.constant 256 : index
|
|
%c32 = arith.constant 32 : index
|
|
scf.for %arg3 = %c0 to %c64 step %c32 {
|
|
scf.for %arg4 = %c0 to %c128 step %c32 {
|
|
%subview = memref.subview %arg2[%arg3, %arg4] [32, 32] [1, 1] : memref<64x128xf32> to !memrefC
|
|
%2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
%3 = vector.transfer_read %subview[%c0, %c16], %0 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
%4 = vector.transfer_read %subview[%c16, %c0], %0 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
%5 = vector.transfer_read %subview[%c16, %c16], %0 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
%6:4 = scf.for %arg5 = %c0 to %c256 step %c64 iter_args(%arg6 = %2, %arg7 = %3, %arg8 = %4, %arg9 = %5) -> (!vecC, !vecC, !vecC, !vecC) {
|
|
%subview_0 = memref.subview %arg0[%arg3, %arg5] [32, 64] [1, 1] : memref<64x256xf8E5M2> to !memrefA
|
|
%subview_1 = memref.subview %arg1[%arg5, %arg4] [64, 32] [1, 1] : memref<256x128xf8E5M2> to !memrefB
|
|
%7 = vector.transfer_read %subview_0[%c0, %c0], %1 {in_bounds = [true, true]} : !memrefA, !vecA
|
|
%8 = vector.transfer_read %subview_0[%c16, %c0], %1 {in_bounds = [true, true]} : !memrefA, !vecA
|
|
%9 = vector.transfer_read %subview_1[%c0, %c0], %1 {in_bounds = [true, true]} : !memrefB, !vecB
|
|
%10 = vector.transfer_read %subview_1[%c0, %c16], %1 {in_bounds = [true, true]} : !memrefB, !vecB
|
|
%11 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%7, %9, %arg6 {unroll_shape = array<i64: 16, 16, 64>} : !vecA, !vecB into !vecC
|
|
%12 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%7, %10, %arg7 {unroll_shape = array<i64: 16, 16, 64>} : !vecA, !vecB into !vecC
|
|
%13 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%8, %9, %arg8 {unroll_shape = array<i64: 16, 16, 64>} : !vecA, !vecB into !vecC
|
|
%14 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%8, %10, %arg9 {unroll_shape = array<i64: 16, 16, 64>} : !vecA, !vecB into !vecC
|
|
scf.yield %11, %12, %13, %14 : !vecC, !vecC, !vecC, !vecC
|
|
}
|
|
vector.transfer_write %6#3, %subview[%c16, %c16] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
vector.transfer_write %6#2, %subview[%c16, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
vector.transfer_write %6#1, %subview[%c0, %c16] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
vector.transfer_write %6#0, %subview[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
}
|
|
}
|
|
%alloc = memref.alloc() : memref<64x128xf32>
|
|
memref.copy %arg2, %alloc : memref<64x128xf32> to memref<64x128xf32>
|
|
return %alloc : memref<64x128xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: @online_packing_f8E5M2_matmul_loop
|
|
// CHECK-COUNT-4: x86.amx.tile_zero : !x86.amx.tile<16x16xf32>
|
|
// CHECK: scf.for {{.*}} -> (!x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>) {
|
|
// CHECK: vector.shuffle{{.*}}[0, 32, 64, 96, 1, 33, 65, 97, 2, 34, 66, 98, 3, 35, 67, 99, 8, 40, 72, 104, 9, 41, 73, 105, 10, 42, 74, 106, 11, 43, 75, 107, 16, 48, 80, 112, 17, 49, 81, 113, 18, 50, 82, 114, 19, 51, 83, 115, 24, 56, 88, 120, 25, 57, 89, 121, 26, 58, 90, 122, 27, 59, 91, 123] : vector<64xf8E5M2>, vector<64xf8E5M2>
|
|
// CHECK-NEXT: vector.shuffle{{.*}}[4, 36, 68, 100, 5, 37, 69, 101, 6, 38, 70, 102, 7, 39, 71, 103, 12, 44, 76, 108, 13, 45, 77, 109, 14, 46, 78, 110, 15, 47, 79, 111, 20, 52, 84, 116, 21, 53, 85, 117, 22, 54, 86, 118, 23, 55, 87, 119, 28, 60, 92, 124, 29, 61, 93, 125, 30, 62, 94, 126, 31, 63, 95, 127] : vector<64xf8E5M2>, vector<64xf8E5M2>
|
|
// CHECK: x86.amx.tile_load
|
|
// CHECK: x86.amx.tile_mulf
|
|
// CHECK: scf.yield {{.*}} !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>
|
|
// CHECK: vector.shuffle{{.*}}[0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23] : vector<16xf32>, vector<16xf32>
|
|
// CHECK-NEXT: vector.shuffle{{.*}}[8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31] : vector<16xf32>, vector<16xf32>
|
|
// CHECK-NOT: scf.for {{.*}} vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>
|
|
// CHECK-NOT: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecAB = vector<1x16x16x2xbf16>
|
|
!vecC = vector<16x16xf32>
|
|
!memrefA = memref<1x32x16x2xbf16, strided<[1024, 32, 2, 1], offset: ?>>
|
|
!memrefB = memref<1x16x32x2xbf16, strided<[1024, 64, 2, 1], offset: ?>>
|
|
|
|
#map = affine_map<(d0, d1, d2, d3, d4) -> (d0, d2, d4, d1)>
|
|
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d4, d3, d1)>
|
|
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d2, d3)>
|
|
|
|
func.func @brgemm_bf16_with_cano(%arg0: memref<16x32x16x2xbf16>, %arg1: memref<16x16x32x2xbf16>, %arg2: memref<32x32xf32>) -> memref<32x32xf32> {
|
|
%0 = ub.poison : f32
|
|
%1 = ub.poison : bf16
|
|
%c0 = arith.constant 0 : index
|
|
%c16 = arith.constant 16 : index
|
|
%c1 = arith.constant 1 : index
|
|
%2 = vector.transfer_read %arg2[%c0, %c0], %0 {in_bounds = [true, true]} : memref<32x32xf32>, !vecC
|
|
%3 = vector.transfer_read %arg2[%c0, %c16], %0 {in_bounds = [true, true]} : memref<32x32xf32>, !vecC
|
|
%4 = vector.transfer_read %arg2[%c16, %c0], %0 {in_bounds = [true, true]} : memref<32x32xf32>, !vecC
|
|
%5 = vector.transfer_read %arg2[%c16, %c16], %0 {in_bounds = [true, true]} : memref<32x32xf32>, !vecC
|
|
|
|
%6:4 = scf.for %arg3 = %c0 to %c16 step %c1 iter_args(%arg4 = %2, %arg5 = %3, %arg6 = %4, %arg7 = %5) -> (!vecC, !vecC, !vecC, !vecC) {
|
|
|
|
%subview = memref.subview %arg0[%arg3, 0, 0, 0] [1, 32, 16, 2] [1, 1, 1, 1] : memref<16x32x16x2xbf16> to !memrefA
|
|
%subview_0 = memref.subview %arg1[%arg3, 0, 0, 0] [1, 16, 32, 2] [1, 1, 1, 1] : memref<16x16x32x2xbf16> to !memrefB
|
|
|
|
%7 = vector.transfer_read %subview[%c0, %c0, %c0, %c0], %1 {in_bounds = [true, true, true, true]} : !memrefA, !vecAB
|
|
%8 = vector.transfer_read %subview[%c0, %c16, %c0, %c0], %1 {in_bounds = [true, true, true, true]} : !memrefA, !vecAB
|
|
%9 = vector.transfer_read %subview_0[%c0, %c0, %c0, %c0], %1 {in_bounds = [true, true, true, true]} : !memrefB, !vecAB
|
|
%10 = vector.transfer_read %subview_0[%c0, %c0, %c16, %c0], %1 {in_bounds = [true, true, true, true]} : !memrefB, !vecAB
|
|
|
|
%11 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%7, %9, %arg4 {unroll_shape = array<i64: 1, 2, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
%12 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%7, %10, %arg5 {unroll_shape = array<i64: 1, 2, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
%13 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%8, %9, %arg6 {unroll_shape = array<i64: 1, 2, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
%14 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%8, %10, %arg7 {unroll_shape = array<i64: 1, 2, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
scf.yield %11, %12, %13, %14 : !vecC, !vecC, !vecC, !vecC
|
|
}
|
|
|
|
vector.transfer_write %6#3, %arg2[%c16, %c16] {in_bounds = [true, true]} : !vecC, memref<32x32xf32>
|
|
vector.transfer_write %6#2, %arg2[%c16, %c0] {in_bounds = [true, true]} : !vecC, memref<32x32xf32>
|
|
vector.transfer_write %6#1, %arg2[%c0, %c16] {in_bounds = [true, true]} : !vecC, memref<32x32xf32>
|
|
vector.transfer_write %6#0, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, memref<32x32xf32>
|
|
%alloc = memref.alloc() : memref<32x32xf32>
|
|
memref.copy %arg2, %alloc : memref<32x32xf32> to memref<32x32xf32>
|
|
return %alloc : memref<32x32xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: @brgemm_bf16_with_cano
|
|
// CHECK-1: scf.for {{.*}} -> (!x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>) {
|
|
// CHECK-4: x86.amx.tile_zero : !x86.amx.tile<16x16xf32>
|
|
// CHECK-4: x86.amx.tile_load
|
|
// CHECK-4: x86.amx.tile_mulf
|
|
// CHECK: scf.yield {{.*}} : !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>, !x86.amx.tile<16x16xf32>
|
|
// CHECK-NOT: scf.for {{.*}} vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>, vector<16x16xf32>
|
|
// CHECK-NOT: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<16x64xi8>
|
|
!vecB = vector<64x16xi8>
|
|
!vecC = vector<16x16xi32>
|
|
!memrefA = memref<32x64xi8, strided<[256, 1], offset: ?>>
|
|
!memrefB = memref<64x32xi8, strided<[128, 1], offset: ?>>
|
|
!memrefC = memref<32x32xi32, strided<[128, 1], offset: ?>>
|
|
|
|
#map = affine_map<(d0, d1, d2) -> (d0, d2)>
|
|
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
|
|
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
func.func @online_packing_int8_matmul_loop(%arg0: memref<64x256xi8>, %arg1: memref<256x128xi8>, %arg2: memref<64x128xi32>) -> memref<64x128xi32> {
|
|
%c16 = arith.constant 16 : index
|
|
%0 = ub.poison : i32
|
|
%1 = ub.poison : i8
|
|
%c0 = arith.constant 0 : index
|
|
%c64 = arith.constant 64 : index
|
|
%c128 = arith.constant 128 : index
|
|
%c256 = arith.constant 256 : index
|
|
%c32 = arith.constant 32 : index
|
|
%bias = arith.constant dense<13> : !vecC
|
|
scf.for %arg3 = %c0 to %c64 step %c32 {
|
|
scf.for %arg4 = %c0 to %c128 step %c32 {
|
|
%subview = memref.subview %arg2[%arg3, %arg4] [32, 32] [1, 1] : memref<64x128xi32> to !memrefC
|
|
%2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
%3 = vector.transfer_read %subview[%c0, %c16], %0 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
%4 = vector.transfer_read %subview[%c16, %c0], %0 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
%5 = vector.transfer_read %subview[%c16, %c16], %0 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
%6:4 = scf.for %arg5 = %c0 to %c256 step %c64 iter_args(%arg6 = %2, %arg7 = %3, %arg8 = %4, %arg9 = %5) -> (!vecC, !vecC, !vecC, !vecC) {
|
|
%subview_0 = memref.subview %arg0[%arg3, %arg5] [32, 64] [1, 1] : memref<64x256xi8> to !memrefA
|
|
%subview_1 = memref.subview %arg1[%arg5, %arg4] [64, 32] [1, 1] : memref<256x128xi8> to !memrefB
|
|
%7 = vector.transfer_read %subview_0[%c0, %c0], %1 {in_bounds = [true, true]} : !memrefA, !vecA
|
|
%8 = vector.transfer_read %subview_0[%c16, %c0], %1 {in_bounds = [true, true]} : !memrefA, !vecA
|
|
%9 = vector.transfer_read %subview_1[%c0, %c0], %1 {in_bounds = [true, true]} : !memrefB, !vecB
|
|
%10 = vector.transfer_read %subview_1[%c0, %c16], %1 {in_bounds = [true, true]} : !memrefB, !vecB
|
|
%11 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%7, %9, %arg6 {unroll_shape = array<i64: 16, 16, 64>} : !vecA, !vecB into !vecC
|
|
%12 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%7, %10, %arg7 {unroll_shape = array<i64: 16, 16, 64>} : !vecA, !vecB into !vecC
|
|
%13 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%8, %9, %arg8 {unroll_shape = array<i64: 16, 16, 64>} : !vecA, !vecB into !vecC
|
|
%14 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%8, %10, %arg9 {unroll_shape = array<i64: 16, 16, 64>} : !vecA, !vecB into !vecC
|
|
scf.yield %11, %12, %13, %14 : !vecC, !vecC, !vecC, !vecC
|
|
}
|
|
|
|
%7 = arith.addi %6#3, %bias : !vecC
|
|
%8 = arith.addi %6#2, %bias : !vecC
|
|
%9 = arith.addi %6#1, %bias : !vecC
|
|
%10 = arith.addi %6#0, %bias : !vecC
|
|
|
|
vector.transfer_write %7, %subview[%c16, %c16] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
vector.transfer_write %8, %subview[%c16, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
vector.transfer_write %9, %subview[%c0, %c16] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
vector.transfer_write %10, %subview[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
}
|
|
}
|
|
%alloc = memref.alloc() : memref<64x128xi32>
|
|
memref.copy %arg2, %alloc : memref<64x128xi32> to memref<64x128xi32>
|
|
return %alloc : memref<64x128xi32>
|
|
}
|
|
|
|
// CHECK-LABEL: @online_packing_int8_matmul_loop
|
|
// CHECK-COUNT-4: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
|
|
// CHECK: scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>) {
|
|
// CHECK: vector.shuffle{{.*}}[0, 32, 64, 96, 1, 33, 65, 97, 2, 34, 66, 98, 3, 35, 67, 99, 8, 40, 72, 104, 9, 41, 73, 105, 10, 42, 74, 106, 11, 43, 75, 107, 16, 48, 80, 112, 17, 49, 81, 113, 18, 50, 82, 114, 19, 51, 83, 115, 24, 56, 88, 120, 25, 57, 89, 121, 26, 58, 90, 122, 27, 59, 91, 123] : vector<64xi8>, vector<64xi8>
|
|
// CHECK-NEXT: vector.shuffle{{.*}}[4, 36, 68, 100, 5, 37, 69, 101, 6, 38, 70, 102, 7, 39, 71, 103, 12, 44, 76, 108, 13, 45, 77, 109, 14, 46, 78, 110, 15, 47, 79, 111, 20, 52, 84, 116, 21, 53, 85, 117, 22, 54, 86, 118, 23, 55, 87, 119, 28, 60, 92, 124, 29, 61, 93, 125, 30, 62, 94, 126, 31, 63, 95, 127] : vector<64xi8>, vector<64xi8>
|
|
// CHECK: x86.amx.tile_load
|
|
// CHECK: x86.amx.tile_muli
|
|
// CHECK: scf.yield {{.*}} !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>
|
|
// CHECK: vector.shuffle{{.*}}[0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23] : vector<16xi32>, vector<16xi32>
|
|
// CHECK-NEXT: vector.shuffle{{.*}}[8, 9, 10, 11, 24, 25, 26, 27, 12, 13, 14, 15, 28, 29, 30, 31] : vector<16xi32>, vector<16xi32>
|
|
// CHECK-COUNT-4: arith.addi
|
|
// CHECK-NOT: scf.for {{.*}} vector<16x16xi32>, vector<16x16xi32>, vector<16x16xi32>, vector<16x16xi32>
|
|
// CHECK-NOT: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<1x16x16x4xi8>
|
|
!vecB = vector<1x16x16x4xi8>
|
|
!vecC = vector<16x16xi32>
|
|
!memrefA = memref<1x32x16x4xi8>
|
|
!memrefB = memref<1x16x32x4xi8>
|
|
!memrefC = memref<32x32xi32>
|
|
#map = affine_map<(d0, d4, d1, d2, d3) -> (d0, d1, d3, d4)>
|
|
#map1 = affine_map<(d0, d4, d1, d2, d3) -> (d0, d3, d2, d4)>
|
|
#map2 = affine_map<(d0, d4, d1, d2, d3) -> (d1, d2)>
|
|
func.func @negative_invalid_vc_kind(
|
|
%arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
|
|
{
|
|
%c0 = arith.constant 0 : index
|
|
%0 = ub.poison : i8
|
|
%32 = ub.poison : i32
|
|
|
|
%1 = vector.transfer_read %arg0[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
|
|
!memrefA, !vecA
|
|
%2 = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
|
|
!memrefB, !vecB
|
|
|
|
%3 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
|
|
%4 = vector.contract {
|
|
indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["reduction", "reduction", "parallel", "parallel", "reduction"],
|
|
kind = #vector.kind<mul>}
|
|
%1, %2, %3 : !vecA, !vecB into !vecC
|
|
|
|
vector.transfer_write %4, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
|
|
return %arg2 : !memrefC
|
|
}
|
|
|
|
// CHECK-LABEL: @negative_invalid_vc_kind
|
|
// CHECK-NOT: x86.amx.tile_load {{.*}} !x86.amx.tile<16x64xi8>
|
|
// CHECK-NOT: x86.amx.tile_muli
|
|
// CHECK-NOT: x86.amx.tile_store {{.*}} !x86.amx.tile<16x16xi32>
|
|
// CHECK: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<1x32x16x2xbf16>
|
|
!vecB = vector<1x16x32x2xbf16>
|
|
!vecC = vector<1x32x32xf32>
|
|
!memrefA = memref<1x32x16x2xbf16>
|
|
!memrefB = memref<1x16x32x2xbf16>
|
|
!memrefC = memref<1x32x32xf32>
|
|
#map = affine_map<(d0, d4, d1, d2, d3) -> (d0, d1, d3, d4)>
|
|
#map1 = affine_map<(d0, d4, d1, d2, d3) -> (d0, d3, d2, d4)>
|
|
#map2 = affine_map<(d0, d4, d1, d2, d3) -> (d0, d1, d2)>
|
|
func.func @negative_wrong_dimensions(
|
|
%arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
|
|
{
|
|
%c0 = arith.constant 0 : index
|
|
%0 = ub.poison : bf16
|
|
%32 = ub.poison : f32
|
|
|
|
%1 = vector.transfer_read %arg0[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
|
|
!memrefA, !vecA
|
|
%2 = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
|
|
!memrefB, !vecB
|
|
|
|
%3 = vector.transfer_read %arg2[%c0, %c0, %c0], %32 {in_bounds = [true, true, true]} : !memrefC, !vecC
|
|
|
|
%4 = vector.contract {
|
|
indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["parallel", "reduction", "parallel", "parallel", "reduction"],
|
|
kind = #vector.kind<add>}
|
|
%1, %2, %3 : !vecA, !vecB into !vecC
|
|
|
|
vector.transfer_write %4, %arg2[%c0, %c0, %c0] {in_bounds = [true, true, true]} : !vecC, !memrefC
|
|
|
|
return %arg2 : !memrefC
|
|
}
|
|
|
|
// CHECK-LABEL: @negative_wrong_dimensions
|
|
// CHECK-NOT: x86.amx.tile_load
|
|
// CHECK-NOT: x86.amx.tile_mulf
|
|
// CHECK-NOT: x86.amx.tile_store
|
|
// CHECK: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<16x16x4xi8>
|
|
!vecB = vector<16x16x4xi8>
|
|
!vecC = vector<16x16xi32>
|
|
!memrefB = memref<16x32x4xi8>
|
|
!memrefC = memref<32x32xi32>
|
|
#map = affine_map<(d4, d1, d2, d3) -> (d1, d3, d4)>
|
|
#map1 = affine_map<(d4, d1, d2, d3) -> (d3, d2, d4)>
|
|
#map2 = affine_map<(d4, d1, d2, d3) -> (d1, d2)>
|
|
func.func @negative_no_memref_source_LHS(
|
|
%arg0: !vecA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
|
|
{
|
|
%c0 = arith.constant 0 : index
|
|
%0 = ub.poison : i8
|
|
%32 = ub.poison : i32
|
|
|
|
%2 = vector.transfer_read %arg1[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} :
|
|
!memrefB, !vecB
|
|
|
|
%3 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
|
|
%4 = vector.contract {
|
|
indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["reduction", "parallel", "parallel", "reduction"],
|
|
kind = #vector.kind<add>}
|
|
%arg0, %2, %3 : !vecA, !vecB into !vecC
|
|
|
|
vector.transfer_write %4, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
|
|
return %arg2 : !memrefC
|
|
}
|
|
|
|
// CHECK-LABEL: @negative_no_memref_source_LHS
|
|
// CHECK-NOT: x86.amx.tile_load {{.*}} !x86.amx.tile<16x64xi8>
|
|
// CHECK-NOT: x86.amx.tile_muli
|
|
// CHECK-NOT: x86.amx.tile_store {{.*}} !x86.amx.tile<16x16xi32>
|
|
// CHECK: vector.contract
|
|
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<1x16x32xbf16>
|
|
!vecB = vector<1x32x32xbf16>
|
|
!vecC = vector<16x32xf32>
|
|
!memrefA = memref<1x32x32xbf16>
|
|
!memrefB = memref<1x32x32xbf16>
|
|
!memrefC = memref<32x32xf32>
|
|
#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
|
|
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
|
|
#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
|
|
func.func @negative_wrong_dimensions_online_packing(
|
|
%arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
|
|
{
|
|
%c0 = arith.constant 0 : index
|
|
%0 = ub.poison : bf16
|
|
%32 = ub.poison : f32
|
|
|
|
%1 = vector.transfer_read %arg0[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} :
|
|
!memrefA, !vecA
|
|
%2 = vector.transfer_read %arg1[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} :
|
|
!memrefB, !vecB
|
|
|
|
%3 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
|
|
%4 = vector.contract {
|
|
indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["reduction", "parallel", "parallel", "reduction"],
|
|
kind = #vector.kind<add>}
|
|
%1, %2, %3 : !vecA, !vecB into !vecC
|
|
|
|
vector.transfer_write %4, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
|
|
return %arg2 : !memrefC
|
|
}
|
|
|
|
// CHECK-LABEL: @negative_wrong_dimensions_online_packing
|
|
// CHECK-NOT: x86.amx.tile_load
|
|
// CHECK-NOT: vector.shuffle
|
|
// CHECK-NOT: x86.amx.tile_mulf
|
|
// CHECK-NOT: x86.amx.tile_store
|
|
// CHECK: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecAB = vector<16x16x4xi8>
|
|
!vecC = vector<16x16xi32>
|
|
!memrefA = memref<16x16x4xi8, strided<[256, 4, 1], offset: ?>>
|
|
!memrefB = memref<16x32x4xi8, strided<[512, 4, 1], offset: ?>>
|
|
!memrefC = memref<16x32xi32, strided<[128, 1], offset: ?>>
|
|
|
|
#map = affine_map<(d0, d1, d2, d3) -> (d1, d3, d0)>
|
|
#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d2, d0)>
|
|
#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
|
|
|
|
func.func @negative_VCs_LHS_src_differ(%arg0: memref<64x64x4xi8>, %arg1: memref<64x128x4xi8>, %arg2: memref<64x128xi32>, %arg13: memref<64x64x4xi8>) {
|
|
%0 = ub.poison : i32
|
|
%1 = ub.poison : i8
|
|
%c0 = arith.constant 0 : index
|
|
%c64 = arith.constant 64 : index
|
|
%c128 = arith.constant 128 : index
|
|
%c16 = arith.constant 16 : index
|
|
%c32 = arith.constant 32 : index
|
|
scf.for %arg3 = %c0 to %c64 step %c16 {
|
|
scf.for %arg4 = %c0 to %c128 step %c32 {
|
|
|
|
%subview = memref.subview %arg2[%arg3, %arg4] [16, 32] [1, 1] :
|
|
memref<64x128xi32> to !memrefC
|
|
%2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
%3 = vector.transfer_read %subview[%c0, %c16], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
|
|
%4:2 = scf.for %arg5 = %c0 to %c64 step %c16 iter_args(%arg6 = %2, %arg7 = %3) -> (!vecC, !vecC) {
|
|
|
|
%subview_0 = memref.subview %arg0[%arg3, %arg5, 0] [16, 16, 4] [1, 1, 1] :
|
|
memref<64x64x4xi8> to !memrefA
|
|
%subview_negative = memref.subview %arg13[%arg3, %arg5, 0] [16, 16, 4] [1, 1, 1] :
|
|
memref<64x64x4xi8> to !memrefA
|
|
%subview_1 = memref.subview %arg1[%arg5, %arg4, 0] [16, 32, 4] [1, 1, 1] :
|
|
memref<64x128x4xi8> to !memrefB
|
|
%5 = vector.transfer_read %subview_0[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
|
|
!memrefA, !vecAB
|
|
%odd_load = vector.transfer_read %subview_negative[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
|
|
!memrefA, !vecAB
|
|
%6 = vector.transfer_read %subview_1[%c0, %c0, %c0], %1 {in_bounds = [true, true, true]} :
|
|
!memrefB, !vecAB
|
|
%7 = vector.transfer_read %subview_1[%c0, %c16, %c0], %1 {in_bounds = [true, true, true]} :
|
|
!memrefB, !vecAB
|
|
|
|
%8 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%5, %6, %arg6 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
%9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%odd_load, %7, %arg7 {unroll_shape = array<i64: 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
|
|
scf.yield %8, %9 : !vecC, !vecC
|
|
}
|
|
|
|
vector.transfer_write %4#1, %subview[%c0, %c16] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
vector.transfer_write %4#0, %subview[%c0, %c0] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
}
|
|
}
|
|
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: @negative_VCs_LHS_src_differ
|
|
// CHECK-NOT: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
|
|
// CHECK-NOT: scf.for {{.*}} -> (!x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>) {
|
|
// CHECK-NOT: x86.amx.tile_load
|
|
// CHECK-NOT: x86.amx.tile_muli
|
|
// CHECK-NOT: scf.yield {{.*}} !x86.amx.tile<16x16xi32>, !x86.amx.tile<16x16xi32>
|
|
// CHECK: scf.for {{.*}} -> (vector<16x16xi32>, vector<16x16xi32>) {
|
|
// CHECK: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<1x16x16x4xi8>
|
|
!vecB = vector<1x16x32x4xi8>
|
|
!vecC = vector<1x16x32xi32>
|
|
!memrefA = memref<1x16x16x4xi8, strided<[16384, 256, 4, 1], offset: ?>>
|
|
!memrefB = memref<1x16x32x4xi8, strided<[32768, 512, 4, 1], offset: ?>>
|
|
!memrefC = memref<1x16x32xi32, strided<[8192, 128, 1], offset: ?>>
|
|
|
|
#map = affine_map<(d0, d1, d2, d3, d4) -> (d0, d2, d4, d1)>
|
|
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d4, d3, d1)>
|
|
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d2, d3)>
|
|
|
|
func.func @negative_wrong_N_dim(%arg0: memref<16x64x64x4xi8>, %arg1: memref<16x64x128x4xi8>, %arg2: memref<16x64x128xi32>) {
|
|
%0 = ub.poison : i32
|
|
%1 = ub.poison : i8
|
|
%c0 = arith.constant 0 : index
|
|
%c64 = arith.constant 64 : index
|
|
%c128 = arith.constant 128 : index
|
|
%c16 = arith.constant 16 : index
|
|
%c32 = arith.constant 32 : index
|
|
%c1 = arith.constant 1 : index
|
|
scf.for %arg3 = %c0 to %c64 step %c16 {
|
|
scf.for %arg4 = %c0 to %c128 step %c32 {
|
|
scf.for %arg5 = %c0 to %c16 step %c1 {
|
|
|
|
%subview = memref.subview %arg2[%arg5, %arg3, %arg4] [1, 16, 32] [1, 1, 1] :
|
|
memref<16x64x128xi32> to !memrefC
|
|
%2 = vector.transfer_read %subview[%c0, %c0, %c0], %0 {in_bounds = [true, true, true]} :
|
|
!memrefC, !vecC
|
|
%3 = scf.for %arg6 = %c0 to %c64 step %c16 iter_args(%arg7 = %2) -> (!vecC) {
|
|
%subview_0 = memref.subview %arg0[%arg5, %arg3, %arg6, 0] [1, 16, 16, 4] [1, 1, 1, 1] :
|
|
memref<16x64x64x4xi8> to !memrefA
|
|
%subview_1 = memref.subview %arg1[%arg5, %arg6, %arg4, 0] [1, 16, 32, 4] [1, 1, 1, 1] :
|
|
memref<16x64x128x4xi8> to !memrefB
|
|
%4 = vector.transfer_read %subview_0[%c0, %c0, %c0, %c0], %1 {in_bounds = [true, true, true, true]} :
|
|
!memrefA, !vecA
|
|
%5 = vector.transfer_read %subview_1[%c0, %c0, %c0, %c0], %1 {in_bounds = [true, true, true, true]} :
|
|
!memrefB, !vecB
|
|
%6 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%4, %5, %arg7 {unroll_shape = array<i64: 1, 4, 16, 32, 16>} : !vecA, !vecB into !vecC
|
|
scf.yield %6 : !vecC
|
|
}
|
|
vector.transfer_write %3, %subview[%c0, %c0, %c0] {in_bounds = [true, true, true]} :
|
|
!vecC, !memrefC
|
|
}
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: @negative_wrong_N_dim
|
|
// CHECK-NOT: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
|
|
// CHECK-NOT: x86.amx.tile_load
|
|
// CHECK-NOT: x86.amx.tile_muli
|
|
// CHECK: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<16x32xbf16>
|
|
!vecB = vector<32x16xbf16>
|
|
!vecC = vector<16x16xf32>
|
|
!memrefA = memref<32x32xbf16, strided<[96, 1], offset: ?>>
|
|
!memrefB = memref<32x32xbf16, strided<[128, 1], offset: ?>>
|
|
!memrefC = memref<32x32xf32, strided<[128, 1], offset: ?>>
|
|
|
|
#map = affine_map<(d1, d2, d3) -> (d1, d3)>
|
|
#map1 = affine_map<(d1, d2, d3) -> (d3, d2)>
|
|
#map2 = affine_map<(d1, d2, d3) -> (d1, d2)>
|
|
|
|
func.func @negative_online_packing_bf16_dynamic_loop_step(%arg0: memref<64x96xbf16>, %arg1: memref<96x128xbf16>, %arg2: memref<64x128xf32>, %kStep: index) -> memref<64x128xf32> {
|
|
%0 = ub.poison : f32
|
|
%1 = ub.poison : bf16
|
|
%c0 = arith.constant 0 : index
|
|
%c64 = arith.constant 64 : index
|
|
%c128 = arith.constant 128 : index
|
|
%c16 = arith.constant 16 : index
|
|
%c96 = arith.constant 96 : index
|
|
%c32 = arith.constant 32 : index
|
|
%c1 = arith.constant 1 : index
|
|
scf.for %arg3 = %c0 to %c64 step %c32 {
|
|
scf.for %arg4 = %c0 to %c128 step %c32 {
|
|
|
|
%subview = memref.subview %arg2[%arg3, %arg4] [32, 32] [1, 1] :
|
|
memref<64x128xf32> to !memrefC
|
|
%2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
%3 = vector.transfer_read %subview[%c0, %c16], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
%4 = vector.transfer_read %subview[%c16, %c0], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
%5 = vector.transfer_read %subview[%c16, %c16], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
|
|
%7:4 = scf.for %arg10 = %c0 to %c32 step %kStep iter_args(%arg11 = %2, %arg12 = %3, %arg13 = %4, %arg14 = %5) -> (!vecC, !vecC, !vecC, !vecC) {
|
|
%subview_0 = memref.subview %arg0[%arg3, %arg10] [32, 32] [1, 1] :
|
|
memref<64x96xbf16> to !memrefA
|
|
%subview_1 = memref.subview %arg1[%arg10, %arg4] [32, 32] [1, 1] :
|
|
memref<96x128xbf16> to !memrefB
|
|
%8 = vector.transfer_read %subview_0[%c0, %c0], %1 {in_bounds = [true, true]} :
|
|
!memrefA, !vecA
|
|
%9 = vector.transfer_read %subview_0[%c16, %c0], %1 {in_bounds = [true, true]} :
|
|
!memrefA, !vecA
|
|
%10 = vector.transfer_read %subview_1[%c0, %c0], %1 {in_bounds = [true, true]} :
|
|
!memrefB, !vecB
|
|
%11 = vector.transfer_read %subview_1[%c0, %c16], %1 {in_bounds = [true, true]} :
|
|
!memrefB, !vecB
|
|
|
|
%12 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%8, %10, %arg11 {unroll_shape = array<i64: 1, 16, 16, 32>} : !vecA, !vecB into !vecC
|
|
%13 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%8, %11, %arg12 {unroll_shape = array<i64: 1, 16, 16, 32>} : !vecA, !vecB into !vecC
|
|
%14 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%9, %10, %arg13 {unroll_shape = array<i64: 1, 16, 16, 32>} : !vecA, !vecB into !vecC
|
|
%15 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%9, %11, %arg14 {unroll_shape = array<i64: 1, 16, 16, 32>} : !vecA, !vecB into !vecC
|
|
|
|
scf.yield %12, %13, %14, %15 : !vecC, !vecC, !vecC, !vecC
|
|
}
|
|
vector.transfer_write %7#3, %subview[%c16, %c16] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
vector.transfer_write %7#2, %subview[%c16, %c0] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
vector.transfer_write %7#1, %subview[%c0, %c16] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
vector.transfer_write %7#0, %subview[%c0, %c0] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
}
|
|
}
|
|
%alloc = memref.alloc() : memref<64x128xf32>
|
|
memref.copy %arg2, %alloc : memref<64x128xf32> to memref<64x128xf32>
|
|
return %alloc : memref<64x128xf32>
|
|
}
|
|
|
|
// CHECK-LABEL: @negative_online_packing_bf16_dynamic_loop_step
|
|
// CHECK-NOT: x86.amx.tile_zero : !x86.amx.tile<16x16xf32>
|
|
// CHECK-NOT: x86.amx.tile_load
|
|
// CHECK-NOT: x86.amx.tile_mulf
|
|
// CHECK: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecAB = vector<1x1x16x16x4xi8>
|
|
!vecC = vector<16x16xi32>
|
|
!memrefA = memref<1x1x16x16x4xi8, strided<[262144, 16384, 256, 4, 1], offset: ?>>
|
|
!memrefB = memref<1x1x16x32x4xi8, strided<[524288, 32768, 512, 4, 1], offset: ?>>
|
|
!memrefC = memref<16x32xi32, strided<[128, 1], offset: ?>>
|
|
|
|
#map = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d3, d5, d2)>
|
|
#map1 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d5, d4, d2)>
|
|
#map2 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4)>
|
|
|
|
func.func @negative_reduction_loop_depth_3(%arg0: memref<2x16x64x64x4xi8>, %arg1: memref<2x16x64x128x4xi8>, %arg2: memref<64x128xi32>) attributes {dlti.target_system_spec = #dlti.target_system_spec<"CPU" = #dlti.target_device_spec<"reg_gemm_unroll" = [16, 16, 16]>>} {
|
|
%0 = ub.poison : i32
|
|
%1 = ub.poison : i8
|
|
%c0 = arith.constant 0 : index
|
|
%c64 = arith.constant 64 : index
|
|
%c128 = arith.constant 128 : index
|
|
%c2 = arith.constant 2 : index
|
|
%c16 = arith.constant 16 : index
|
|
%c32 = arith.constant 32 : index
|
|
%c1 = arith.constant 1 : index
|
|
scf.for %arg3 = %c0 to %c64 step %c16 {
|
|
scf.for %arg4 = %c0 to %c128 step %c32 {
|
|
%subview = memref.subview %arg2[%arg3, %arg4] [16, 32] [1, 1] :
|
|
memref<64x128xi32> to !memrefC
|
|
%2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
%3 = vector.transfer_read %subview[%c0, %c16], %0 {in_bounds = [true, true]} :
|
|
!memrefC, !vecC
|
|
|
|
%4:2 = scf.for %arg5 = %c0 to %c2 step %c1 iter_args(%arg6 = %2, %arg7 = %3) -> (!vecC, !vecC) {
|
|
%5:2 = scf.for %arg8 = %c0 to %c16 step %c1 iter_args(%arg9 = %arg6, %arg10 = %arg7) -> (!vecC, !vecC) {
|
|
%6:2 = scf.for %arg11 = %c0 to %c64 step %c16 iter_args(%arg12 = %arg9, %arg13 = %arg10) -> (!vecC, !vecC) {
|
|
%subview_0 = memref.subview %arg0[%arg5, %arg8, %arg3, %arg11, 0] [1, 1, 16, 16, 4] [1, 1, 1, 1, 1] :
|
|
memref<2x16x64x64x4xi8> to !memrefA
|
|
%subview_1 = memref.subview %arg1[%arg5, %arg8, %arg11, %arg4, 0] [1, 1, 16, 32, 4] [1, 1, 1, 1, 1] :
|
|
memref<2x16x64x128x4xi8> to !memrefB
|
|
%7 = vector.transfer_read %subview_0[%c0, %c0, %c0, %c0, %c0], %1 {in_bounds = [true, true, true, true, true]} :
|
|
!memrefA, !vecAB
|
|
%8 = vector.transfer_read %subview_1[%c0, %c0, %c0, %c0, %c0], %1 {in_bounds = [true, true, true, true, true]} :
|
|
!memrefB, !vecAB
|
|
%9 = vector.transfer_read %subview_1[%c0, %c0, %c0, %c16, %c0], %1 {in_bounds = [true, true, true, true, true]} :
|
|
!memrefB, !vecAB
|
|
|
|
%10 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "reduction", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%7, %8, %arg12 {unroll_shape = array<i64: 1, 1, 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
%11 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["reduction", "reduction", "reduction", "parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%7, %9, %arg13 {unroll_shape = array<i64: 1, 1, 4, 16, 16, 16>} : !vecAB, !vecAB into !vecC
|
|
scf.yield %10, %11 : !vecC, !vecC
|
|
}
|
|
scf.yield %6#0, %6#1 : !vecC, !vecC
|
|
}
|
|
scf.yield %5#0, %5#1 : !vecC, !vecC
|
|
}
|
|
|
|
vector.transfer_write %4#1, %subview[%c0, %c16] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
vector.transfer_write %4#0, %subview[%c0, %c0] {in_bounds = [true, true]} :
|
|
!vecC, !memrefC
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
|
|
// CHECK-LABEL: @negative_reduction_loop_depth_3
|
|
// CHECK-NOT: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
|
|
// CHECK-NOT: x86.amx.tile_load
|
|
// CHECK-NOT: x86.amx.tile_muli
|
|
// CHECK: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<1x16x16x2xf16>
|
|
!vecB = vector<1x16x16x2xf16>
|
|
!vecC = vector<16x16xf32>
|
|
!memrefA = memref<1x32x16x2xf16>
|
|
!memrefB = memref<1x16x32x2xf16>
|
|
!memrefC = memref<32x32xf32>
|
|
#map = affine_map<(d0, d4, d1, d2, d3) -> (d0, d1, d3, d4)>
|
|
#map1 = affine_map<(d0, d4, d1, d2, d3) -> (d0, d3, d2, d4)>
|
|
#map2 = affine_map<(d0, d4, d1, d2, d3) -> (d1, d2)>
|
|
func.func @negative_wrong_type_f16(
|
|
%arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
|
|
{
|
|
%c0 = arith.constant 0 : index
|
|
%0 = ub.poison : f16
|
|
%32 = ub.poison : f32
|
|
|
|
%1 = vector.transfer_read %arg0[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
|
|
!memrefA, !vecA
|
|
%2 = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
|
|
!memrefB, !vecB
|
|
|
|
%3 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
|
|
%4 = vector.contract {
|
|
indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["reduction", "reduction", "parallel", "parallel", "reduction"],
|
|
kind = #vector.kind<add>}
|
|
%1, %2, %3 : !vecA, !vecB into !vecC
|
|
|
|
vector.transfer_write %4, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
|
|
return %arg2 : !memrefC
|
|
}
|
|
|
|
// CHECK-LABEL: @negative_wrong_type_f16
|
|
// CHECK-NOT: x86.amx.tile_load
|
|
// CHECK-NOT: x86.amx.tile_mulf
|
|
// CHECK-NOT: x86.amx.tile_store
|
|
// CHECK: vector.contract
|
|
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<1x16x16x2xbf16>
|
|
!vecB = vector<1x16x16x2xbf16>
|
|
!vecC = vector<16x16xbf16>
|
|
!memrefA = memref<1x32x16x2xbf16>
|
|
!memrefB = memref<1x16x32x2xbf16>
|
|
!memrefC = memref<32x32xbf16>
|
|
#map = affine_map<(d0, d4, d1, d2, d3) -> (d0, d1, d3, d4)>
|
|
#map1 = affine_map<(d0, d4, d1, d2, d3) -> (d0, d3, d2, d4)>
|
|
#map2 = affine_map<(d0, d4, d1, d2, d3) -> (d1, d2)>
|
|
func.func @negative_wrong_acc_type_bf16(
|
|
%arg0: !memrefA, %arg1: !memrefB, %arg2: !memrefC) -> !memrefC
|
|
{
|
|
%c0 = arith.constant 0 : index
|
|
%0 = ub.poison : bf16
|
|
%32 = ub.poison : bf16
|
|
|
|
%1 = vector.transfer_read %arg0[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
|
|
!memrefA, !vecA
|
|
%2 = vector.transfer_read %arg1[%c0, %c0, %c0, %c0], %0 {in_bounds = [true, true, true, true]} :
|
|
!memrefB, !vecB
|
|
|
|
%3 = vector.transfer_read %arg2[%c0, %c0], %32 {in_bounds = [true, true]} : !memrefC, !vecC
|
|
|
|
%4 = vector.contract {
|
|
indexing_maps = [#map, #map1, #map2],
|
|
iterator_types = ["reduction", "reduction", "parallel", "parallel", "reduction"],
|
|
kind = #vector.kind<add>}
|
|
%1, %2, %3 : !vecA, !vecB into !vecC
|
|
|
|
vector.transfer_write %4, %arg2[%c0, %c0] {in_bounds = [true, true]} : !vecC, !memrefC
|
|
|
|
return %arg2 : !memrefC
|
|
}
|
|
|
|
// CHECK-LABEL: @negative_wrong_acc_type_bf16
|
|
// CHECK-NOT: x86.amx.tile_load
|
|
// CHECK-NOT: x86.amx.tile_mulf
|
|
// CHECK-NOT: x86.amx.tile_store
|
|
// CHECK: vector.contract
|
|
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<16x64xi8>
|
|
!vecB = vector<64x16xi8>
|
|
!vecC = vector<16x16xi32>
|
|
!memrefA = memref<16x64xi8, strided<[256, 1], offset: ?>>
|
|
!memrefB = memref<64x32xi8, strided<[128, 1], offset: ?>>
|
|
!memrefC = memref<16x32xi32, strided<[128, 1], offset: ?>>
|
|
|
|
#map = affine_map<(d0, d1, d2) -> (d0, d2)>
|
|
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
|
|
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
|
|
func.func @negative_vc_wrong_order_no_pair(%arg0: memref<64x256xi8>, %arg1: memref<256x128xi8>, %arg2: memref<64x128xi32>) -> memref<64x128xi32> {
|
|
%0 = ub.poison : i32
|
|
%1 = ub.poison : i8
|
|
%c0 = arith.constant 0 : index
|
|
%c64 = arith.constant 64 : index
|
|
%c128 = arith.constant 128 : index
|
|
%c256 = arith.constant 256 : index
|
|
%c16 = arith.constant 16 : index
|
|
%c32 = arith.constant 32 : index
|
|
scf.for %arg3 = %c0 to %c64 step %c16 {
|
|
scf.for %arg4 = %c0 to %c128 step %c32 {
|
|
%subview = memref.subview %arg2[%arg3, %arg4] [16, 32] [1, 1]
|
|
: memref<64x128xi32> to !memrefC
|
|
%2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]}
|
|
: !memrefC, !vecC
|
|
%3 = vector.transfer_read %subview[%c0, %c16], %0 {in_bounds = [true, true]}
|
|
: !memrefC, !vecC
|
|
%4:2 = scf.for %arg5 = %c0 to %c256 step %c64 iter_args(%arg6 = %2, %arg7 = %3) -> (!vecC, !vecC) {
|
|
%subview_0 = memref.subview %arg0[%arg3, %arg5] [16, 64] [1, 1]
|
|
: memref<64x256xi8> to !memrefA
|
|
%subview_1 = memref.subview %arg1[%arg5, %arg4] [64, 32] [1, 1]
|
|
: memref<256x128xi8> to !memrefB
|
|
%5 = vector.transfer_read %subview_0[%c0, %c0], %1 {in_bounds = [true, true]}
|
|
: !memrefA, !vecA
|
|
%6 = vector.transfer_read %subview_1[%c0, %c0], %1 {in_bounds = [true, true]}
|
|
: !memrefB, !vecB
|
|
%7 = vector.transfer_read %subview_1[%c0, %c16], %1 {in_bounds = [true, true]}
|
|
: !memrefB, !vecB
|
|
%8 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%5, %7, %arg7 {unroll_shape = array<i64: 16, 16, 64>} : !vecA, !vecB into !vecC
|
|
%9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%5, %6, %arg6 {unroll_shape = array<i64: 16, 16, 64>} : !vecA, !vecB into !vecC
|
|
scf.yield %9, %8 : !vecC, !vecC
|
|
}
|
|
vector.transfer_write %4#1, %subview[%c0, %c16] {in_bounds = [true, true]}
|
|
: !vecC, !memrefC
|
|
vector.transfer_write %4#0, %subview[%c0, %c0] {in_bounds = [true, true]}
|
|
: !vecC, !memrefC
|
|
}
|
|
}
|
|
%alloc = memref.alloc() : memref<64x128xi32>
|
|
memref.copy %arg2, %alloc : memref<64x128xi32> to memref<64x128xi32>
|
|
return %alloc : memref<64x128xi32>
|
|
}
|
|
|
|
|
|
// CHECK-LABEL: @negative_vc_wrong_order_no_pair
|
|
// CHECK-NOT: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
|
|
// CHECK-NOT: x86.amx.tile_load
|
|
// CHECK-NOT: x86.amx.tile_muli
|
|
// CHECK: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
!vecA = vector<16x64xi8>
|
|
!vecB = vector<64x16xi8>
|
|
!vecC = vector<16x16xi32>
|
|
!memrefA = memref<32x64xi8, strided<[256, 1], offset: ?>>
|
|
!memrefB = memref<64x16xi8, strided<[128, 1], offset: ?>>
|
|
!memrefC = memref<32x16xi32, strided<[128, 1], offset: ?>>
|
|
|
|
#map = affine_map<(d0, d1, d2) -> (d0, d2)>
|
|
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
|
|
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
|
|
func.func @negative_vc_no_pair(%arg0: memref<64x256xi8>, %arg1: memref<256x128xi8>, %arg2: memref<64x128xi32>) -> memref<64x128xi32> {
|
|
%0 = ub.poison : i32
|
|
%1 = ub.poison : i8
|
|
%c0 = arith.constant 0 : index
|
|
%c64 = arith.constant 64 : index
|
|
%c128 = arith.constant 128 : index
|
|
%c256 = arith.constant 256 : index
|
|
%c32 = arith.constant 32 : index
|
|
%c16 = arith.constant 16 : index
|
|
scf.for %arg3 = %c0 to %c64 step %c32 {
|
|
scf.for %arg4 = %c0 to %c128 step %c16 {
|
|
%subview = memref.subview %arg2[%arg3, %arg4] [32, 16] [1, 1] : memref<64x128xi32> to !memrefC
|
|
%2 = vector.transfer_read %subview[%c0, %c0], %0 {in_bounds = [true, true]}
|
|
: !memrefC, !vecC
|
|
%3 = vector.transfer_read %subview[%c16, %c0], %0 {in_bounds = [true, true]}
|
|
: !memrefC, !vecC
|
|
%4:2 = scf.for %arg5 = %c0 to %c256 step %c64 iter_args(%arg6 = %2, %arg7 = %3) -> (!vecC, !vecC) {
|
|
%subview_0 = memref.subview %arg0[%arg3, %arg5] [32, 64] [1, 1]
|
|
: memref<64x256xi8> to !memrefA
|
|
%subview_1 = memref.subview %arg1[%arg5, %arg4] [64, 16] [1, 1]
|
|
: memref<256x128xi8> to !memrefB
|
|
%5 = vector.transfer_read %subview_0[%c0, %c0], %1 {in_bounds = [true, true]}
|
|
: !memrefA, !vecA
|
|
%6 = vector.transfer_read %subview_0[%c16, %c0], %1 {in_bounds = [true, true]}
|
|
: !memrefA, !vecA
|
|
%7 = vector.transfer_read %subview_1[%c0, %c0], %1 {in_bounds = [true, true]}
|
|
: !memrefB, !vecB
|
|
%8 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%5, %7, %arg6 {unroll_shape = array<i64: 16, 16, 64>} : !vecA, !vecB into !vecC
|
|
%9 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types =
|
|
["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
|
|
%6, %7, %arg7 {unroll_shape = array<i64: 16, 16, 64>} : !vecA, !vecB into !vecC
|
|
scf.yield %8, %9 : !vecC, !vecC
|
|
}
|
|
vector.transfer_write %4#1, %subview[%c16, %c0] {in_bounds = [true, true]}
|
|
: !vecC, !memrefC
|
|
vector.transfer_write %4#0, %subview[%c0, %c0] {in_bounds = [true, true]}
|
|
: !vecC, !memrefC
|
|
}
|
|
}
|
|
%alloc = memref.alloc() : memref<64x128xi32>
|
|
memref.copy %arg2, %alloc : memref<64x128xi32> to memref<64x128xi32>
|
|
return %alloc : memref<64x128xi32>
|
|
}
|
|
|
|
// CHECK-LABEL: @negative_vc_no_pair
|
|
// CHECK-NOT: x86.amx.tile_zero : !x86.amx.tile<16x16xi32>
|
|
// CHECK-NOT: x86.amx.tile_load
|
|
// CHECK-NOT: x86.amx.tile_muli
|
|
// CHECK: vector.contract
|
|
|
|
module attributes {transform.with_named_sequence} {
|
|
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
|
%func = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
|
|
transform.apply_patterns to %func {
|
|
transform.apply_patterns.x86.vector_contract_to_amx_dot_product
|
|
} : !transform.any_op
|
|
transform.yield
|
|
}
|
|
}
|