mlir/test/Conversion/MemRefToSPIRV/map-storage-class.mlir

   1 // RUN: mlir-opt -split-input-file -allow-unregistered-dialect -map-memref-spirv-storage-class='client-api=vulkan' -verify-diagnostics %s -o - | FileCheck %s --check-prefix=VULKAN
   2 // RUN: mlir-opt -split-input-file -allow-unregistered-dialect -map-memref-spirv-storage-class='client-api=opencl' -verify-diagnostics %s -o - | FileCheck %s --check-prefix=OPENCL
   3
   4 // Vulkan Mappings:
   5 //   0 -> StorageBuffer
   6 //   1 -> Generic
   7 //   2 -> [null]
   8 //   3 -> Workgroup
   9 //   4 -> Uniform
  10
  11 // OpenCL Mappings:
  12 //   0 -> CrossWorkgroup
  13 //   1 -> Generic
  14 //   2 -> [null]
  15 //   3 -> Workgroup
  16 //   4 -> UniformConstant
  17
  18 // VULKAN-LABEL: func @operand_result
  19 // OPENCL-LABEL: func @operand_result
  20 func.func @operand_result() {
  21   // VULKAN: memref<f32, #spirv.storage_class<StorageBuffer>>
  22   // OPENCL: memref<f32, #spirv.storage_class<CrossWorkgroup>>
  23   %0 = "dialect.memref_producer"() : () -> (memref<f32>)
  24   // VULKAN: memref<4xi32, #spirv.storage_class<Generic>>
  25   // OPENCL: memref<4xi32, #spirv.storage_class<Generic>>
  26   %1 = "dialect.memref_producer"() : () -> (memref<4xi32, 1>)
  27   // VULKAN: memref<?x4xf16, #spirv.storage_class<Workgroup>>
  28   // OPENCL: memref<?x4xf16, #spirv.storage_class<Workgroup>>
  29   %2 = "dialect.memref_producer"() : () -> (memref<?x4xf16, 3>)
  30   // VULKAN: memref<*xf16, #spirv.storage_class<Uniform>>
  31   // OPENCL: memref<*xf16, #spirv.storage_class<UniformConstant>>
  32   %3 = "dialect.memref_producer"() : () -> (memref<*xf16, 4>)
  33
  34
  35   "dialect.memref_consumer"(%0) : (memref<f32>) -> ()
  36   // VULKAN: memref<4xi32, #spirv.storage_class<Generic>>
  37   // OPENCL: memref<4xi32, #spirv.storage_class<Generic>>
  38   "dialect.memref_consumer"(%1) : (memref<4xi32, 1>) -> ()
  39   // VULKAN: memref<?x4xf16, #spirv.storage_class<Workgroup>>
  40   // OPENCL: memref<?x4xf16, #spirv.storage_class<Workgroup>>
  41   "dialect.memref_consumer"(%2) : (memref<?x4xf16, 3>) -> ()
  42   // VULKAN: memref<*xf16, #spirv.storage_class<Uniform>>
  43   // OPENCL: memref<*xf16, #spirv.storage_class<UniformConstant>>
  44   "dialect.memref_consumer"(%3) : (memref<*xf16, 4>) -> ()
  45
  46   return
  47 }
  48
  49 // -----
  50
  51 // VULKAN-LABEL: func @type_attribute
  52 // OPENCL-LABEL: func @type_attribute
  53 func.func @type_attribute() {
  54   // VULKAN: attr = memref<i32, #spirv.storage_class<Generic>>
  55   // OPENCL: attr = memref<i32, #spirv.storage_class<Generic>>
  56   "dialect.memref_producer"() { attr = memref<i32, 1> } : () -> ()
  57   return
  58 }
  59
  60 // -----
  61
  62 // VULKAN-LABEL: func.func @function_io
  63 // OPENCL-LABEL: func.func @function_io
  64 func.func @function_io
  65   // VULKAN-SAME: (%{{.+}}: memref<f64, #spirv.storage_class<Generic>>, %{{.+}}: memref<4xi32, #spirv.storage_class<Workgroup>>)
  66   // OPENCL-SAME: (%{{.+}}: memref<f64, #spirv.storage_class<Generic>>, %{{.+}}: memref<4xi32, #spirv.storage_class<Workgroup>>)
  67   (%arg0: memref<f64, 1>, %arg1: memref<4xi32, 3>)
  68   // VULKAN-SAME: -> (memref<f64, #spirv.storage_class<Generic>>, memref<4xi32, #spirv.storage_class<Workgroup>>)
  69   // OPENCL-SAME: -> (memref<f64, #spirv.storage_class<Generic>>, memref<4xi32, #spirv.storage_class<Workgroup>>)
  70   -> (memref<f64, 1>, memref<4xi32, 3>) {
  71   return %arg0, %arg1: memref<f64, 1>, memref<4xi32, 3>
  72 }
  73
  74 // -----
  75
  76 gpu.module @kernel {
  77 // VULKAN-LABEL: gpu.func @function_io
  78 // OPENCL-LABEL: gpu.func @function_io
  79 // VULKAN-SAME: memref<8xi32, #spirv.storage_class<StorageBuffer>>
  80 // OPENCL-SAME: memref<8xi32, #spirv.storage_class<CrossWorkgroup>>
  81 gpu.func @function_io(%arg0 : memref<8xi32>) kernel { gpu.return }
  82 }
  83
  84 // -----
  85
  86 // VULKAN-LABEL: func.func @region
  87 // OPENCL-LABEL: func.func @region
  88 func.func @region(%cond: i1, %arg0: memref<f32, 1>) {
  89   scf.if %cond {
  90     //      VULKAN: "dialect.memref_consumer"(%{{.+}}) {attr = memref<i64, #spirv.storage_class<Workgroup>>}
  91     //      OPENCL: "dialect.memref_consumer"(%{{.+}}) {attr = memref<i64, #spirv.storage_class<Workgroup>>}
  92     // VULKAN-SAME: (memref<f32, #spirv.storage_class<Generic>>) -> memref<f32, #spirv.storage_class<Generic>>
  93     // OPENCL-SAME: (memref<f32, #spirv.storage_class<Generic>>) -> memref<f32, #spirv.storage_class<Generic>>
  94     %0 = "dialect.memref_consumer"(%arg0) { attr = memref<i64, 3> } : (memref<f32, 1>) -> (memref<f32, 1>)
  95   }
  96   return
  97 }
  98
  99 // -----
 100
 101 // VULKAN-LABEL: func @non_memref_types
 102 // OPENCL-LABEL: func @non_memref_types
 103 func.func @non_memref_types(%arg: f32) -> f32 {
 104   // VULKAN: "dialect.op"(%{{.+}}) {attr = 16 : i64} : (f32) -> f32
 105   // OPENCL: "dialect.op"(%{{.+}}) {attr = 16 : i64} : (f32) -> f32
 106   %0 = "dialect.op"(%arg) { attr = 16 } : (f32) -> (f32)
 107   return %0 : f32
 108 }
 109
 110 // -----
 111
 112 func.func @missing_mapping() {
 113   // expected-error @+1 {{failed to legalize}}
 114   %0 = "dialect.memref_producer"() : () -> (memref<f32, 2>)
 115   return
 116 }
 117
 118 // -----
 119
 120 /// Checks memory maps to OpenCL mapping if Kernel capability is enabled.
 121 module attributes { spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Kernel], []>, #spirv.resource_limits<>> } {
 122 func.func @operand_result() {
 123   // CHECK: memref<f32, #spirv.storage_class<CrossWorkgroup>>
 124   %0 = "dialect.memref_producer"() : () -> (memref<f32>)
 125   // CHECK: memref<4xi32, #spirv.storage_class<Generic>>
 126   %1 = "dialect.memref_producer"() : () -> (memref<4xi32, 1>)
 127   // CHECK: memref<?x4xf16, #spirv.storage_class<Workgroup>>
 128   %2 = "dialect.memref_producer"() : () -> (memref<?x4xf16, 3>)
 129   // CHECK: memref<*xf16, #spirv.storage_class<UniformConstant>>
 130   %3 = "dialect.memref_producer"() : () -> (memref<*xf16, 4>)
 131
 132
 133   "dialect.memref_consumer"(%0) : (memref<f32>) -> ()
 134   // CHECK: memref<4xi32, #spirv.storage_class<Generic>>
 135   "dialect.memref_consumer"(%1) : (memref<4xi32, 1>) -> ()
 136   // CHECK: memref<?x4xf16, #spirv.storage_class<Workgroup>>
 137   "dialect.memref_consumer"(%2) : (memref<?x4xf16, 3>) -> ()
 138   // CHECK: memref<*xf16, #spirv.storage_class<UniformConstant>>
 139   "dialect.memref_consumer"(%3) : (memref<*xf16, 4>) -> ()
 140
 141   return
 142 }
 143 }
 144
 145 // -----
 146
 147 /// Checks memory maps to Vulkan mapping if Shader capability is enabled.
 148 module attributes { spirv.target_env = #spirv.target_env<#spirv.vce<v1.0, [Shader], []>, #spirv.resource_limits<>> } {
 149 func.func @operand_result() {
 150   // CHECK: memref<f32, #spirv.storage_class<StorageBuffer>>
 151   %0 = "dialect.memref_producer"() : () -> (memref<f32>)
 152   // CHECK: memref<4xi32, #spirv.storage_class<Generic>>
 153   %1 = "dialect.memref_producer"() : () -> (memref<4xi32, 1>)
 154   // CHECK: memref<?x4xf16, #spirv.storage_class<Workgroup>>
 155   %2 = "dialect.memref_producer"() : () -> (memref<?x4xf16, 3>)
 156   // CHECK: memref<*xf16, #spirv.storage_class<Uniform>>
 157   %3 = "dialect.memref_producer"() : () -> (memref<*xf16, 4>)
 158
 159
 160   "dialect.memref_consumer"(%0) : (memref<f32>) -> ()
 161   // CHECK: memref<4xi32, #spirv.storage_class<Generic>>
 162   "dialect.memref_consumer"(%1) : (memref<4xi32, 1>) -> ()
 163   // CHECK: memref<?x4xf16, #spirv.storage_class<Workgroup>>
 164   "dialect.memref_consumer"(%2) : (memref<?x4xf16, 3>) -> ()
 165   // CHECK: memref<*xf16, #spirv.storage_class<Uniform>>
 166   "dialect.memref_consumer"(%3) : (memref<*xf16, 4>) -> ()
 167   return
 168 }
 169 }