mlir/test/Target/SPIRV/joint-matrix-ops.mlir

   1 // RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip -split-input-file %s | FileCheck %s
   2
   3 spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [JointMatrixINTEL], [SPV_INTEL_joint_matrix]> {
   4   // CHECK-LABEL: @joint_matrix_load
   5   spirv.func @joint_matrix_load(%ptr : !spirv.ptr<i32, Workgroup>, %stride : i32) "None" {
   6     // CHECK: {{%.*}} = spirv.INTEL.JointMatrixLoad <Subgroup> <RowMajor> {{%.*}}, {{%.*}} : (!spirv.ptr<i32, Workgroup>, i32) -> !spirv.jointmatrix<16x8xi32, RowMajor, Workgroup>
   7     %0 = spirv.INTEL.JointMatrixLoad <Subgroup> <RowMajor> %ptr, %stride : (!spirv.ptr<i32, Workgroup>, i32) -> !spirv.jointmatrix<16x8xi32, RowMajor, Workgroup>
   8     spirv.Return
   9   }
  10
  11   // CHECK-LABEL: @joint_matrix_load_memaccess
  12   spirv.func @joint_matrix_load_memaccess(%ptr : !spirv.ptr<i32, Workgroup>, %stride : i32) "None" {
  13     // CHECK: {{%.*}} = spirv.INTEL.JointMatrixLoad <Subgroup> <RowMajor> {{%.*}}, {{%.*}} {memory_access = #spirv.memory_access<Volatile>} : (!spirv.ptr<i32, Workgroup>, i32) -> !spirv.jointmatrix<8x16xi32, RowMajor, Subgroup>
  14     %0 = spirv.INTEL.JointMatrixLoad <Subgroup> <RowMajor> %ptr, %stride {memory_access = #spirv.memory_access<Volatile>} : (!spirv.ptr<i32, Workgroup>, i32) -> !spirv.jointmatrix<8x16xi32, RowMajor, Subgroup>
  15     spirv.Return
  16   }
  17
  18   // CHECK-LABEL: @joint_matrix_store
  19   spirv.func @joint_matrix_store(%ptr : !spirv.ptr<i32, Workgroup>, %stride : i32, %m : !spirv.jointmatrix<16x8xi32, RowMajor, Workgroup>) "None" {
  20     // CHECK: spirv.INTEL.JointMatrixStore <Subgroup> <RowMajor> {{%.*}}, {{%.*}}, {{%.*}} : (!spirv.ptr<i32, Workgroup>, !spirv.jointmatrix<16x8xi32, RowMajor, Workgroup>, i32)
  21     spirv.INTEL.JointMatrixStore <Subgroup> <RowMajor> %ptr, %m, %stride : (!spirv.ptr<i32, Workgroup>, !spirv.jointmatrix<16x8xi32, RowMajor, Workgroup>, i32)
  22     spirv.Return
  23   }
  24
  25   // CHECK-LABEL: @joint_matrix_store_memaccess
  26   spirv.func @joint_matrix_store_memaccess(%ptr : !spirv.ptr<i32, Workgroup>, %m : !spirv.jointmatrix<8x16xi32, RowMajor, Subgroup>, %stride : i32) "None" {
  27     // CHECK: spirv.INTEL.JointMatrixStore <Subgroup> <RowMajor> {{%.*}}, {{%.*}}, {{%.*}} {memory_access = #spirv.memory_access<Volatile>} : (!spirv.ptr<i32, Workgroup>, !spirv.jointmatrix<8x16xi32, RowMajor, Subgroup>, i32)
  28     spirv.INTEL.JointMatrixStore <Subgroup> <RowMajor> %ptr, %m, %stride {memory_access = #spirv.memory_access<Volatile>} : (!spirv.ptr<i32, Workgroup>, !spirv.jointmatrix<8x16xi32, RowMajor, Subgroup>, i32)
  29     spirv.Return
  30   }
  31
  32   // CHECK-LABEL: @joint_matrix_length
  33   spirv.func @joint_matrix_length() -> i32 "None" {
  34     // CHECK: {{%.*}} = spirv.INTEL.JointMatrixWorkItemLength : !spirv.jointmatrix<8x16xi32, RowMajor, Subgroup>
  35     %0 = spirv.INTEL.JointMatrixWorkItemLength : !spirv.jointmatrix<8x16xi32, RowMajor, Subgroup>
  36     spirv.ReturnValue %0 : i32
  37   }
  38
  39   // CHECK-LABEL: @joint_matrix_muladd
  40   spirv.func @joint_matrix_muladd(%a : !spirv.jointmatrix<8x16xi32, RowMajor, Subgroup>, %b : !spirv.jointmatrix<16x8xi32, RowMajor, Subgroup>, %c : !spirv.jointmatrix<8x8xi32, RowMajor, Subgroup>) "None" {
  41     // CHECK: {{%.*}} = spirv.INTEL.JointMatrixMad <Subgroup> {{%.*}}, {{%.*}}, {{%.*}}  : !spirv.jointmatrix<8x16xi32, RowMajor, Subgroup>, !spirv.jointmatrix<16x8xi32, RowMajor, Subgroup> -> !spirv.jointmatrix<8x8xi32, RowMajor, Subgroup>
  42     %r = spirv.INTEL.JointMatrixMad <Subgroup> %a, %b, %c : !spirv.jointmatrix<8x16xi32, RowMajor, Subgroup>, !spirv.jointmatrix<16x8xi32, RowMajor, Subgroup> -> !spirv.jointmatrix<8x8xi32, RowMajor, Subgroup>
  43     spirv.Return
  44   }
  45 }