mlir/test/Target/SPIRV/cooperative-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, [CooperativeMatrixNV], [SPV_NV_cooperative_matrix]> {
   4   // CHECK-LABEL: @cooperative_matrix_load
   5   spirv.func @cooperative_matrix_load(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32, %b : i1) "None" {
   6     // CHECK: {{%.*}} = spirv.NV.CooperativeMatrixLoad {{%.*}}, {{%.*}}, {{%.*}} : !spirv.ptr<i32, StorageBuffer> as !spirv.coopmatrix<16x8xi32, Workgroup>
   7     %0 = spirv.NV.CooperativeMatrixLoad %ptr, %stride, %b : !spirv.ptr<i32, StorageBuffer> as !spirv.coopmatrix<16x8xi32, Workgroup>
   8     spirv.Return
   9   }
  10
  11   // CHECK-LABEL: @cooperative_matrix_load_memaccess
  12   spirv.func @cooperative_matrix_load_memaccess(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32, %b : i1) "None" {
  13     // CHECK: {{%.*}} = spirv.NV.CooperativeMatrixLoad {{%.*}}, {{%.*}}, {{%.*}} ["Volatile"] : !spirv.ptr<i32, StorageBuffer> as !spirv.coopmatrix<8x16xi32, Subgroup>
  14     %0 = spirv.NV.CooperativeMatrixLoad %ptr, %stride, %b ["Volatile"] : !spirv.ptr<i32, StorageBuffer> as !spirv.coopmatrix<8x16xi32, Subgroup>
  15     spirv.Return
  16   }
  17
  18   // CHECK-LABEL: @cooperative_matrix_store
  19   spirv.func @cooperative_matrix_store(%ptr : !spirv.ptr<i32, StorageBuffer>, %stride : i32, %m : !spirv.coopmatrix<16x8xi32, Workgroup>, %b : i1) "None" {
  20     // CHECK: spirv.NV.CooperativeMatrixStore {{%.*}}, {{%.*}}, {{%.*}} : !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<16x8xi32, Workgroup>
  21     spirv.NV.CooperativeMatrixStore %ptr, %m, %stride, %b : !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<16x8xi32, Workgroup>
  22     spirv.Return
  23   }
  24
  25   // CHECK-LABEL: @cooperative_matrix_store_memaccess
  26   spirv.func @cooperative_matrix_store_memaccess(%ptr : !spirv.ptr<i32, StorageBuffer>, %m : !spirv.coopmatrix<8x16xi32, Subgroup>, %stride : i32, %b : i1) "None" {
  27     // CHECK: spirv.NV.CooperativeMatrixStore {{%.*}}, {{%.*}}, {{%.*}} ["Volatile"] : !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<8x16xi32, Subgroup>
  28     spirv.NV.CooperativeMatrixStore %ptr, %m, %stride, %b ["Volatile"] : !spirv.ptr<i32, StorageBuffer>, !spirv.coopmatrix<8x16xi32, Subgroup>
  29     spirv.Return
  30   }
  31
  32   // CHECK-LABEL: @cooperative_matrix_length
  33   spirv.func @cooperative_matrix_length() -> i32 "None" {
  34     // CHECK: {{%.*}} = spirv.NV.CooperativeMatrixLength : !spirv.coopmatrix<8x16xi32, Subgroup>
  35     %0 = spirv.NV.CooperativeMatrixLength : !spirv.coopmatrix<8x16xi32, Subgroup>
  36     spirv.ReturnValue %0 : i32
  37   }
  38
  39   // CHECK-LABEL: @cooperative_matrix_muladd
  40   spirv.func @cooperative_matrix_muladd(%a : !spirv.coopmatrix<8x16xi32, Subgroup>, %b : !spirv.coopmatrix<16x8xi32, Subgroup>, %c : !spirv.coopmatrix<8x8xi32, Subgroup>) "None" {
  41     // CHECK: {{%.*}} = spirv.NV.CooperativeMatrixMulAdd {{%.*}}, {{%.*}}, {{%.*}}  : !spirv.coopmatrix<8x16xi32, Subgroup>, !spirv.coopmatrix<16x8xi32, Subgroup> -> !spirv.coopmatrix<8x8xi32, Subgroup>
  42     %r = spirv.NV.CooperativeMatrixMulAdd %a, %b, %c : !spirv.coopmatrix<8x16xi32, Subgroup>, !spirv.coopmatrix<16x8xi32, Subgroup> -> !spirv.coopmatrix<8x8xi32, Subgroup>
  43     spirv.Return
  44   }
  45
  46   // CHECK-LABEL: @cooperative_matrix_add
  47   spirv.func @cooperative_matrix_add(%a : !spirv.coopmatrix<8x16xi32, Subgroup>, %b : !spirv.coopmatrix<8x16xi32, Subgroup>) "None" {
  48     // CHECK: {{%.*}} = spirv.IAdd {{%.*}}, {{%.*}} : !spirv.coopmatrix<8x16xi32, Subgroup>
  49     %r = spirv.IAdd %a, %b : !spirv.coopmatrix<8x16xi32, Subgroup>
  50     spirv.Return
  51   }
  52
  53   // CHECK-LABEL: @cooperative_matrix_sub
  54   spirv.func @cooperative_matrix_sub(%a : !spirv.coopmatrix<8x16xi32, Subgroup>, %b : !spirv.coopmatrix<8x16xi32, Subgroup>) "None" {
  55     // CHECK: {{%.*}} = spirv.ISub {{%.*}}, {{%.*}} : !spirv.coopmatrix<8x16xi32, Subgroup>
  56     %r = spirv.ISub %a, %b : !spirv.coopmatrix<8x16xi32, Subgroup>
  57     spirv.Return
  58   }
  59
  60   // CHECK-LABEL: @cooperative_matrix_sdiv
  61   spirv.func @cooperative_matrix_sdiv(%a : !spirv.coopmatrix<8x16xi32, Subgroup>, %b : !spirv.coopmatrix<8x16xi32, Subgroup>) "None" {
  62     // CHECK: {{%.*}} = spirv.SDiv {{%.*}}, {{%.*}} : !spirv.coopmatrix<8x16xi32, Subgroup>
  63     %r = spirv.SDiv %a, %b : !spirv.coopmatrix<8x16xi32, Subgroup>
  64     spirv.Return
  65   }
  66
  67   // CHECK-LABEL: @cooperative_matrix_udiv
  68   spirv.func @cooperative_matrix_udiv(%a : !spirv.coopmatrix<8x16xi32, Subgroup>, %b : !spirv.coopmatrix<8x16xi32, Subgroup>) "None" {
  69     // CHECK: {{%.*}} = spirv.UDiv {{%.*}}, {{%.*}} : !spirv.coopmatrix<8x16xi32, Subgroup>
  70     %r = spirv.UDiv %a, %b : !spirv.coopmatrix<8x16xi32, Subgroup>
  71     spirv.Return
  72   }
  73
  74   // CHECK-LABEL: @cooperative_matrix_fadd
  75   spirv.func @cooperative_matrix_fadd(%a : !spirv.coopmatrix<8x16xf32, Subgroup>, %b : !spirv.coopmatrix<8x16xf32, Subgroup>) "None" {
  76     // CHECK: {{%.*}} = spirv.FAdd {{%.*}}, {{%.*}} : !spirv.coopmatrix<8x16xf32, Subgroup>
  77     %r = spirv.FAdd %a, %b : !spirv.coopmatrix<8x16xf32, Subgroup>
  78     spirv.Return
  79   }
  80
  81   // CHECK-LABEL: @cooperative_matrix_fsub
  82   spirv.func @cooperative_matrix_fsub(%a : !spirv.coopmatrix<8x16xf32, Subgroup>, %b : !spirv.coopmatrix<8x16xf32, Subgroup>) "None" {
  83     // CHECK: {{%.*}} = spirv.FSub {{%.*}}, {{%.*}} : !spirv.coopmatrix<8x16xf32, Subgroup>
  84     %r = spirv.FSub %a, %b : !spirv.coopmatrix<8x16xf32, Subgroup>
  85     spirv.Return
  86   }
  87
  88   // CHECK-LABEL: @cooperative_matrix_fdiv
  89   spirv.func @cooperative_matrix_fdiv(%a : !spirv.coopmatrix<8x16xf32, Subgroup>, %b : !spirv.coopmatrix<8x16xf32, Subgroup>) "None" {
  90     // CHECK: {{%.*}} = spirv.FDiv {{%.*}}, {{%.*}} : !spirv.coopmatrix<8x16xf32, Subgroup>
  91     %r = spirv.FDiv %a, %b : !spirv.coopmatrix<8x16xf32, Subgroup>
  92     spirv.Return
  93   }
  94
  95   // CHECK-LABEL: @cooperative_matrix_access_chain
  96   spirv.func @cooperative_matrix_access_chain(%a : !spirv.ptr<!spirv.coopmatrix<8x16xf32, Subgroup>, Function>) -> !spirv.ptr<f32, Function> "None" {
  97     %0 = spirv.Constant 0: i32
  98     // CHECK: {{%.*}} = spirv.AccessChain {{%.*}}[{{%.*}}] : !spirv.ptr<!spirv.coopmatrix<8x16xf32, Subgroup>, Function>, i32
  99     %1 = spirv.AccessChain %a[%0] : !spirv.ptr<!spirv.coopmatrix<8x16xf32, Subgroup>, Function>, i32
 100     spirv.ReturnValue %1 : !spirv.ptr<f32, Function>
 101   }
 102 }