[mlir][spirv] Query target environment for mapping memory space

[raikonenfnu]

Summary

Checks spirv::TargetEnv from op to see if it contains either Kernel or Shader capabilities.
If it does, then it will set the memory space mapping accordingly.

[raikonenfnu]

Updated commit message

mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRVPass.h"
 
 #include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
+#include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Debug.h"
 
 namespace mlir {
 #define GEN_PASS_DEF_MAPMEMREFSTORAGECLASS
LogicalResult MapMemRefStorageClassPass::initializeOptions(StringRef options) {
 
   return success();
 }
 
 void MapMemRefStorageClassPass::runOnOperation() {
   MLIRContext *context = &getContext();
   Operation *op = getOperation();
 
+  if (spirv::TargetEnvAttr attr = spirv::lookupTargetEnv(op)) {
+    spirv::TargetEnv targetEnv(attr);
+    if (targetEnv.allows(spirv::Capability::Kernel)) {
+      memorySpaceMap = spirv::mapMemorySpaceToOpenCLStorageClass;
+    } else if (targetEnv.allows(spirv::Capability::Shader)) {
+      memorySpaceMap = spirv::mapMemorySpaceToVulkanStorageClass;
+    }
+  }
+
   auto target = spirv::getMemorySpaceToStorageClassTarget(*context);
   spirv::MemorySpaceToStorageClassConverter converter(memorySpaceMap);
 
   RewritePatternSet patterns(context);
   spirv::populateMemorySpaceToStorageClassPatterns(converter, patterns);
 
   if (failed(applyFullConversion(op, *target, std::move(patterns))))
     return signalPassFailure();
 }
 
 std::unique_ptr<OperationPass<>> mlir::createMapMemRefStorageClassPass() {
   return std::make_unique<MapMemRefStorageClassPass>();
 }

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

 
 // -----
 
 func.func @missing_mapping() {
   // expected-error @+1 {{failed to legalize}}
   %0 = "dialect.memref_producer"() : () -> (memref<f32, 2>)
   return
 }
+
+// -----
+
+/// Checks memory maps to OpenCL mapping if Kernel capability is enabled.
+module attributes { spv.target_env = #spv.target_env<#spv.vce<v1.0, [Kernel], []>, #spv.resource_limits<>> } {
+func.func @operand_result() {
+  // CHECK: memref<f32, #spv.storage_class<CrossWorkgroup>>
+  %0 = "dialect.memref_producer"() : () -> (memref<f32>)
+  // CHECK: memref<4xi32, #spv.storage_class<Generic>>
+  %1 = "dialect.memref_producer"() : () -> (memref<4xi32, 1>)
+  // CHECK: memref<?x4xf16, #spv.storage_class<Workgroup>>
+  %2 = "dialect.memref_producer"() : () -> (memref<?x4xf16, 3>)
+  // CHECK: memref<*xf16, #spv.storage_class<UniformConstant>>
+  %3 = "dialect.memref_producer"() : () -> (memref<*xf16, 4>)
+
+
+  "dialect.memref_consumer"(%0) : (memref<f32>) -> ()
+  // CHECK: memref<4xi32, #spv.storage_class<Generic>>
+  "dialect.memref_consumer"(%1) : (memref<4xi32, 1>) -> ()
+  // CHECK: memref<?x4xf16, #spv.storage_class<Workgroup>>
+  "dialect.memref_consumer"(%2) : (memref<?x4xf16, 3>) -> ()
+  // CHECK: memref<*xf16, #spv.storage_class<UniformConstant>>
+  "dialect.memref_consumer"(%3) : (memref<*xf16, 4>) -> ()
+
+  return
+}
+}
+
+// -----
+
+/// Checks memory maps to Vulkan mapping if Shader capability is enabled.
+module attributes { spv.target_env = #spv.target_env<#spv.vce<v1.0, [Shader], []>, #spv.resource_limits<>> } {
+func.func @operand_result() {
+  // CHECK: memref<f32, #spv.storage_class<StorageBuffer>>
+  %0 = "dialect.memref_producer"() : () -> (memref<f32>)
+  // CHECK: memref<4xi32, #spv.storage_class<Generic>>
+  %1 = "dialect.memref_producer"() : () -> (memref<4xi32, 1>)
+  // CHECK: memref<?x4xf16, #spv.storage_class<Workgroup>>
+  %2 = "dialect.memref_producer"() : () -> (memref<?x4xf16, 3>)
+  // CHECK: memref<*xf16, #spv.storage_class<Uniform>>
+  %3 = "dialect.memref_producer"() : () -> (memref<*xf16, 4>)
+
+
+  "dialect.memref_consumer"(%0) : (memref<f32>) -> ()
+  // CHECK: memref<4xi32, #spv.storage_class<Generic>>
+  "dialect.memref_consumer"(%1) : (memref<4xi32, 1>) -> ()
+  // CHECK: memref<?x4xf16, #spv.storage_class<Workgroup>>
+  "dialect.memref_consumer"(%2) : (memref<?x4xf16, 3>) -> ()
+  // CHECK: memref<*xf16, #spv.storage_class<Uniform>>
+  "dialect.memref_consumer"(%3) : (memref<*xf16, 4>) -> ()
+  return
+}
+}
 No newline at end of file