130 lines · plain
1// RUN: mlir-translate --no-implicit-module --test-spirv-roundtrip \2// RUN: --split-input-file %s | FileCheck %s3 4// RUN: %if spirv-tools %{ rm -rf %t %}5// RUN: %if spirv-tools %{ mkdir %t %}6// RUN: %if spirv-tools %{ mlir-translate --no-implicit-module --serialize-spirv --split-input-file --spirv-save-validation-files-with-prefix=%t/module %s %}7// RUN: %if spirv-tools %{ spirv-val %t %}8 9spirv.module Logical Vulkan requires10 #spirv.vce<v1.5, [Shader, Int8, Int16, Int64, Linkage, CooperativeMatrixKHR, VulkanMemoryModel],11 [SPV_KHR_storage_buffer_storage_class, SPV_KHR_cooperative_matrix]> {12 13 // CHECK-LABEL: @cooperative_matrix_length14 spirv.func @cooperative_matrix_length() "None" {15 // CHECK: {{%.+}} = spirv.KHR.CooperativeMatrixLength : !spirv.coopmatrix<2x2xi32, Subgroup, MatrixB>16 %0 = spirv.KHR.CooperativeMatrixLength : !spirv.coopmatrix<2x2xi32, Subgroup, MatrixB>17 spirv.Return18 }19 20 // CHECK-LABEL: @cooperative_matrix_load_121 spirv.func @cooperative_matrix_load_1(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32) "None" {22 // CHECK: {{%.+}} = spirv.KHR.CooperativeMatrixLoad {{%.*}}, {{%.*}}, <RowMajor>23 // CHECK-SAME: : !spirv.ptr<i32, StorageBuffer>, i32 -> !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>24 %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, <RowMajor> :25 !spirv.ptr<i32, StorageBuffer>, i32 -> !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>26 spirv.Return27 }28 29 // CHECK-LABEL: @cooperative_matrix_load_230 spirv.func @cooperative_matrix_load_2(%ptr : !spirv.ptr<f32, StorageBuffer>, %stride : i64) "None" {31 // CHECK: {{%.+}} = spirv.KHR.CooperativeMatrixLoad {{%.*}}, {{%.*}}, <ColumnMajor>, <Volatile>32 // CHECK-SAME: : !spirv.ptr<f32, StorageBuffer>, i64 -> !spirv.coopmatrix<8x16xf32, Subgroup, MatrixAcc>33 %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, <ColumnMajor>, <Volatile> :34 !spirv.ptr<f32, StorageBuffer>, i64 -> !spirv.coopmatrix<8x16xf32, Subgroup, MatrixAcc>35 spirv.Return36 }37 38 // CHECK-LABEL: @cooperative_matrix_load_339 spirv.func @cooperative_matrix_load_3(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32) "None" {40 // CHECK: {{%.+}} = spirv.KHR.CooperativeMatrixLoad {{%.*}}, {{%.*}}, <RowMajor>, <Aligned>, 1641 // CHECK-SAME: : !spirv.ptr<i32, StorageBuffer>, i32 -> !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>42 %0 = spirv.KHR.CooperativeMatrixLoad %ptr, %stride, <RowMajor>, <Aligned>, 16 :43 !spirv.ptr<i32, StorageBuffer>, i32 -> !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>44 spirv.Return45 }46 47 // CHECK-LABEL: @cooperative_matrix_store_148 spirv.func @cooperative_matrix_store_1(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32,49 %m : !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>) "None" {50 // CHECK: spirv.KHR.CooperativeMatrixStore {{%.*}}, {{%.*}}, <RowMajor>51 // CHECK-SAME: : !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>, i3252 spirv.KHR.CooperativeMatrixStore %ptr, %m, %stride, <RowMajor> :53 !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>, i3254 55 // CHECK: spirv.KHR.CooperativeMatrixStore {{%.*}}, {{%.*}}, <RowMajor>, <Aligned>, 1656 // CHECK-SAME: : !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>, i3257 spirv.KHR.CooperativeMatrixStore %ptr, %m, %stride, <RowMajor>, <Aligned>, 16 :58 !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>, i3259 60 // CHECK-NEXT: spirv.KHR.CooperativeMatrixStore {{%.*}}, {{%.*}}, <RowMajor>, <Volatile|Nontemporal>61 spirv.KHR.CooperativeMatrixStore %ptr, %m, %stride, <RowMajor>, <Volatile|Nontemporal> :62 !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<16x8xi32, Workgroup, MatrixA>, i3263 spirv.Return64 }65 66 // CHECK-LABEL: @cooperative_matrix_store_267 spirv.func @cooperative_matrix_store_2(%ptr : !spirv.ptr<f32, Workgroup>, %stride : i64,68 %m : !spirv.coopmatrix<4x8xf32, Subgroup, MatrixB>) "None" {69 // CHECK: spirv.KHR.CooperativeMatrixStore {{%.*}}, {{%.*}}, <ColumnMajor>, <Nontemporal>70 // CHECK-SAME: : !spirv.ptr<f32, Workgroup>, !spirv.coopmatrix<4x8xf32, Subgroup, MatrixB>, i6471 spirv.KHR.CooperativeMatrixStore %ptr, %m, %stride, <ColumnMajor>, <Nontemporal> :72 !spirv.ptr<f32, Workgroup>, !spirv.coopmatrix<4x8xf32, Subgroup, MatrixB>, i6473 spirv.Return74 }75 76 // CHECK-LABEL: @cooperative_matrix_muladd77 spirv.func @cooperative_matrix_muladd_1(%a : !spirv.coopmatrix<8x16xi8, Subgroup, MatrixA>,78 %b : !spirv.coopmatrix<16x8xi16, Subgroup, MatrixB>,79 %c : !spirv.coopmatrix<8x8xi32, Subgroup, MatrixAcc>) "None" {80 // CHECK: {{%.+}} = spirv.KHR.CooperativeMatrixMulAdd {{%.*}}, {{%.*}}, {{%.*}} :81 // CHECK-SAME: !spirv.coopmatrix<8x16xi8, Subgroup, MatrixA>,82 // CHECK-SAME: !spirv.coopmatrix<16x8xi16, Subgroup, MatrixB>83 // CHECK-SAME: -> !spirv.coopmatrix<8x8xi32, Subgroup, MatrixAcc>84 %p = spirv.KHR.CooperativeMatrixMulAdd %a, %b, %c : !spirv.coopmatrix<8x16xi8, Subgroup, MatrixA>,85 !spirv.coopmatrix<16x8xi16, Subgroup, MatrixB>86 -> !spirv.coopmatrix<8x8xi32, Subgroup, MatrixAcc>87 88 // CHECK-NEXT: {{%.+}} = spirv.KHR.CooperativeMatrixMulAdd {{%.*}}, {{%.*}}, {{%.*}}, <ASigned> :89 // CHECK-SAME: !spirv.coopmatrix<8x16xi8, Subgroup, MatrixA>,90 // CHECK-SAME: !spirv.coopmatrix<16x8xi16, Subgroup, MatrixB>91 // CHECK-SAME: -> !spirv.coopmatrix<8x8xi32, Subgroup, MatrixAcc>92 %q = spirv.KHR.CooperativeMatrixMulAdd %a, %b, %c,93 <ASigned> : !spirv.coopmatrix<8x16xi8, Subgroup, MatrixA>,94 !spirv.coopmatrix<16x8xi16, Subgroup, MatrixB>95 -> !spirv.coopmatrix<8x8xi32, Subgroup, MatrixAcc>96 97 // CHECK-NEXT: {{%.+}} = spirv.KHR.CooperativeMatrixMulAdd {{%.*}}, {{%.*}}, {{%.*}}, <BSigned> :98 %r = spirv.KHR.CooperativeMatrixMulAdd %a, %b, %c,99 <BSigned> : !spirv.coopmatrix<8x16xi8, Subgroup, MatrixA>,100 !spirv.coopmatrix<16x8xi16, Subgroup, MatrixB>101 -> !spirv.coopmatrix<8x8xi32, Subgroup, MatrixAcc>102 103 // CHECK-NEXT: {{%.+}} = spirv.KHR.CooperativeMatrixMulAdd104 // CHECK-SAME: <ASigned|BSigned|ResultSigned|AccSat> :105 %s = spirv.KHR.CooperativeMatrixMulAdd %a, %b, %c,106 <ASigned|BSigned|ResultSigned|AccSat> :107 !spirv.coopmatrix<8x16xi8, Subgroup, MatrixA>,108 !spirv.coopmatrix<16x8xi16, Subgroup, MatrixB>109 -> !spirv.coopmatrix<8x8xi32, Subgroup, MatrixAcc>110 111 spirv.Return112 }113 114 // CHECK-LABEL: @cooperative_matrix_muladd115 spirv.func @cooperative_matrix_muladd_2(%a : !spirv.coopmatrix<8x8xf32, Workgroup, MatrixA>,116 %b : !spirv.coopmatrix<8x8xf32, Workgroup, MatrixB>,117 %c : !spirv.coopmatrix<8x8xf32, Workgroup, MatrixAcc>) "None" {118 // CHECK: {{%.+}} = spirv.KHR.CooperativeMatrixMulAdd {{%.*}}, {{%.*}}, {{%.*}} :119 // CHECK-SAME: !spirv.coopmatrix<8x8xf32, Workgroup, MatrixA>,120 // CHECK-SAME: !spirv.coopmatrix<8x8xf32, Workgroup, MatrixB>121 // CHECK-SAME: -> !spirv.coopmatrix<8x8xf32, Workgroup, MatrixAcc>122 %r = spirv.KHR.CooperativeMatrixMulAdd %a, %b, %c : !spirv.coopmatrix<8x8xf32, Workgroup, MatrixA>,123 !spirv.coopmatrix<8x8xf32, Workgroup, MatrixB>124 -> !spirv.coopmatrix<8x8xf32, Workgroup, MatrixAcc>125 126 spirv.Return127 }128 129}130