Files
TheRock/compiler/amd-llvm/mlir/test/Conversion/VectorToSPIRV/vector-reduction-to-spirv-dot-prod.mlir
T
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
Flatten subprojects into a single base monorepo
2026-09-02 16:34:47 -07:00

193 lines
9.0 KiB
MLIR

// RUN: mlir-opt --split-input-file --verify-diagnostics \
// RUN: --test-vector-reduction-to-spirv-dot-prod %s -o - | FileCheck %s
// Positive tests.
// CHECK-LABEL: func.func @to_sdot
// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>)
// CHECK-NEXT: [[DOT:%.+]] = spirv.SDot [[ARG0]], [[ARG1]] : vector<4xi8> -> i32
// CHECK-NEXT: return [[DOT]] : i32
func.func @to_sdot(%arg0: vector<4xi8>, %arg1: vector<4xi8>) -> i32 {
%lhs = arith.extsi %arg0 : vector<4xi8> to vector<4xi32>
%rhs = arith.extsi %arg1 : vector<4xi8> to vector<4xi32>
%mul = arith.muli %lhs, %rhs : vector<4xi32>
%red = vector.reduction <add>, %mul : vector<4xi32> into i32
return %red : i32
}
// CHECK-LABEL: func.func @to_sdot_acc
// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>, [[ACC:%.+]]: i32)
// CHECK-NEXT: [[DOT:%.+]] = spirv.SDotAccSat [[ARG0]], [[ARG1]], [[ACC]] : vector<4xi8> -> i32
// CHECK-NEXT: return [[DOT]] : i32
func.func @to_sdot_acc(%arg0: vector<4xi8>, %arg1: vector<4xi8>, %acc: i32) -> i32 {
%lhs = arith.extsi %arg0 : vector<4xi8> to vector<4xi32>
%rhs = arith.extsi %arg1 : vector<4xi8> to vector<4xi32>
%mul = arith.muli %lhs, %rhs : vector<4xi32>
%red = vector.reduction <add>, %mul, %acc : vector<4xi32> into i32
return %red : i32
}
// CHECK-LABEL: func.func @to_sdot_i64
// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>)
// CHECK-NEXT: [[DOT:%.+]] = spirv.SDot [[ARG0]], [[ARG1]] : vector<4xi8> -> i64
// CHECK-NEXT: return [[DOT]] : i64
func.func @to_sdot_i64(%arg0: vector<4xi8>, %arg1: vector<4xi8>) -> i64 {
%lhs = arith.extsi %arg0 : vector<4xi8> to vector<4xi64>
%rhs = arith.extsi %arg1 : vector<4xi8> to vector<4xi64>
%mul = arith.muli %lhs, %rhs : vector<4xi64>
%red = vector.reduction <add>, %mul : vector<4xi64> into i64
return %red : i64
}
// CHECK-LABEL: func.func @to_sdot_acc_i64
// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>, [[ACC:%.+]]: i64)
// CHECK-NEXT: [[DOT:%.+]] = spirv.SDotAccSat [[ARG0]], [[ARG1]], [[ACC]] : vector<4xi8> -> i64
// CHECK-NEXT: return [[DOT]] : i64
func.func @to_sdot_acc_i64(%arg0: vector<4xi8>, %arg1: vector<4xi8>, %acc: i64) -> i64 {
%lhs = arith.extsi %arg0 : vector<4xi8> to vector<4xi64>
%rhs = arith.extsi %arg1 : vector<4xi8> to vector<4xi64>
%mul = arith.muli %lhs, %rhs : vector<4xi64>
%red = vector.reduction <add>, %mul, %acc : vector<4xi64> into i64
return %red : i64
}
// CHECK-LABEL: func.func @to_udot
// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>)
// CHECK-NEXT: [[DOT:%.+]] = spirv.UDot [[ARG0]], [[ARG1]] : vector<4xi8> -> i32
// CHECK-NEXT: return [[DOT]] : i32
func.func @to_udot(%arg0: vector<4xi8>, %arg1: vector<4xi8>) -> i32 {
%lhs = arith.extui %arg0 : vector<4xi8> to vector<4xi32>
%rhs = arith.extui %arg1 : vector<4xi8> to vector<4xi32>
%mul = arith.muli %lhs, %rhs : vector<4xi32>
%red = vector.reduction <add>, %mul : vector<4xi32> into i32
return %red : i32
}
// CHECK-LABEL: func.func @to_udot_acc
// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>, [[ACC:%.+]]: i32)
// CHECK-NEXT: [[DOT:%.+]] = spirv.UDotAccSat [[ARG0]], [[ARG1]], [[ACC]] : vector<4xi8> -> i32
// CHECK-NEXT: return [[DOT]] : i32
func.func @to_udot_acc(%arg0: vector<4xi8>, %arg1: vector<4xi8>, %acc: i32) -> i32 {
%lhs = arith.extui %arg0 : vector<4xi8> to vector<4xi32>
%rhs = arith.extui %arg1 : vector<4xi8> to vector<4xi32>
%mul = arith.muli %lhs, %rhs : vector<4xi32>
%red = vector.reduction <add>, %mul, %acc : vector<4xi32> into i32
return %red : i32
}
// CHECK-LABEL: func.func @to_signed_unsigned_dot
// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>)
// CHECK-NEXT: [[DOT:%.+]] = spirv.SUDot [[ARG0]], [[ARG1]] : vector<4xi8> -> i32
// CHECK-NEXT: return [[DOT]] : i32
func.func @to_signed_unsigned_dot(%arg0: vector<4xi8>, %arg1: vector<4xi8>) -> i32 {
%lhs = arith.extsi %arg0 : vector<4xi8> to vector<4xi32>
%rhs = arith.extui %arg1 : vector<4xi8> to vector<4xi32>
%mul = arith.muli %lhs, %rhs : vector<4xi32>
%red = vector.reduction <add>, %mul : vector<4xi32> into i32
return %red : i32
}
// CHECK-LABEL: func.func @to_signed_unsigned_dot_acc
// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>, [[ACC:%.+]]: i32)
// CHECK-NEXT: [[DOT:%.+]] = spirv.SUDotAccSat [[ARG0]], [[ARG1]], [[ACC]] : vector<4xi8> -> i32
// CHECK-NEXT: return [[DOT]] : i32
func.func @to_signed_unsigned_dot_acc(%arg0: vector<4xi8>, %arg1: vector<4xi8>, %acc: i32) -> i32 {
%lhs = arith.extsi %arg0 : vector<4xi8> to vector<4xi32>
%rhs = arith.extui %arg1 : vector<4xi8> to vector<4xi32>
%mul = arith.muli %lhs, %rhs : vector<4xi32>
%red = vector.reduction <add>, %mul, %acc : vector<4xi32> into i32
return %red : i32
}
// CHECK-LABEL: func.func @to_unsigned_signed_dot
// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>)
// CHECK-NEXT: [[DOT:%.+]] = spirv.SUDot [[ARG1]], [[ARG0]] : vector<4xi8> -> i32
// CHECK-NEXT: return [[DOT]] : i32
func.func @to_unsigned_signed_dot(%arg0: vector<4xi8>, %arg1: vector<4xi8>) -> i32 {
%lhs = arith.extui %arg0 : vector<4xi8> to vector<4xi32>
%rhs = arith.extsi %arg1 : vector<4xi8> to vector<4xi32>
%mul = arith.muli %lhs, %rhs : vector<4xi32>
%red = vector.reduction <add>, %mul : vector<4xi32> into i32
return %red : i32
}
// CHECK-LABEL: func.func @to_unsigned_signed_dot_acc
// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>, [[ACC:%.+]]: i32)
// CHECK-NEXT: [[DOT:%.+]] = spirv.SUDotAccSat [[ARG1]], [[ARG0]], [[ACC]] : vector<4xi8> -> i32
// CHECK-NEXT: return [[DOT]] : i32
func.func @to_unsigned_signed_dot_acc(%arg0: vector<4xi8>, %arg1: vector<4xi8>, %acc: i32) -> i32 {
%lhs = arith.extui %arg0 : vector<4xi8> to vector<4xi32>
%rhs = arith.extsi %arg1 : vector<4xi8> to vector<4xi32>
%mul = arith.muli %lhs, %rhs : vector<4xi32>
%red = vector.reduction <add>, %mul, %acc : vector<4xi32> into i32
return %red : i32
}
// CHECK-LABEL: func.func @to_sdot_vector3
// CHECK-SAME: (%[[ARG0:.+]]: vector<3xi8>, %[[ARG1:.+]]: vector<3xi8>)
// CHECK: %[[ZERO:.+]] = spirv.Constant 0 : i8
// CHECK: %[[LHS:.+]] = spirv.CompositeConstruct %[[ARG0]], %[[ZERO]] : (vector<3xi8>, i8) -> vector<4xi8>
// CHECK: %[[RHS:.+]] = spirv.CompositeConstruct %[[ARG1]], %[[ZERO]] : (vector<3xi8>, i8) -> vector<4xi8>
// CHECK: %[[SDOT:.+]] = spirv.SDot %[[LHS]], %[[RHS]] : vector<4xi8> -> i32
// CHECK: return %[[SDOT]]
func.func @to_sdot_vector3(%arg0: vector<3xi8>, %arg1: vector<3xi8>) -> i32 {
%lhs = arith.extsi %arg0 : vector<3xi8> to vector<3xi32>
%rhs = arith.extsi %arg1 : vector<3xi8> to vector<3xi32>
%mul = arith.muli %lhs, %rhs : vector<3xi32>
%red = vector.reduction <add>, %mul : vector<3xi32> into i32
return %red : i32
}
// -----
// Negative tests.
// CHECK-LABEL: func.func @too_short
// CHECK-SAME: ([[ARG0:%.+]]: vector<2xi8>, [[ARG1:%.+]]: vector<2xi8>)
// CHECK: [[RED:%.+]] = vector.reduction
// CHECK-NEXT: return [[RED]] : i32
func.func @too_short(%arg0: vector<2xi8>, %arg1: vector<2xi8>) -> i32 {
%lhs = arith.extsi %arg0 : vector<2xi8> to vector<2xi32>
%rhs = arith.extsi %arg1 : vector<2xi8> to vector<2xi32>
%mul = arith.muli %lhs, %rhs : vector<2xi32>
%red = vector.reduction <add>, %mul : vector<2xi32> into i32
return %red : i32
}
// CHECK-LABEL: func.func @too_long
// CHECK-SAME: ([[ARG0:%.+]]: vector<6xi8>, [[ARG1:%.+]]: vector<6xi8>)
// CHECK: [[RED:%.+]] = vector.reduction
// CHECK-NEXT: return [[RED]] : i32
func.func @too_long(%arg0: vector<6xi8>, %arg1: vector<6xi8>) -> i32 {
%lhs = arith.extsi %arg0 : vector<6xi8> to vector<6xi32>
%rhs = arith.extsi %arg1 : vector<6xi8> to vector<6xi32>
%mul = arith.muli %lhs, %rhs : vector<6xi32>
%red = vector.reduction <add>, %mul : vector<6xi32> into i32
return %red : i32
}
// CHECK-LABEL: func.func @wrong_reduction_kind
// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>)
// CHECK: [[RED:%.+]] = vector.reduction <mul>
// CHECK-NEXT: return [[RED]] : i32
func.func @wrong_reduction_kind(%arg0: vector<4xi8>, %arg1: vector<4xi8>) -> i32 {
%lhs = arith.extsi %arg0 : vector<4xi8> to vector<4xi32>
%rhs = arith.extsi %arg1 : vector<4xi8> to vector<4xi32>
%mul = arith.muli %lhs, %rhs : vector<4xi32>
%red = vector.reduction <mul>, %mul : vector<4xi32> into i32
return %red : i32
}
// CHECK-LABEL: func.func @wrong_arith_op
// CHECK-SAME: ([[ARG0:%.+]]: vector<4xi8>, [[ARG1:%.+]]: vector<4xi8>)
// CHECK: [[ADD:%.+]] = arith.addi
// CHECK: [[RED:%.+]] = vector.reduction <mul>, [[ADD]]
// CHECK-NEXT: return [[RED]] : i32
func.func @wrong_arith_op(%arg0: vector<4xi8>, %arg1: vector<4xi8>) -> i32 {
%lhs = arith.extsi %arg0 : vector<4xi8> to vector<4xi32>
%rhs = arith.extsi %arg1 : vector<4xi8> to vector<4xi32>
%add = arith.addi %lhs, %rhs : vector<4xi32>
%red = vector.reduction <mul>, %add : vector<4xi32> into i32
return %red : i32
}