diff --git a/mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h b/mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h --- a/mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h +++ b/mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h @@ -23,6 +23,8 @@ namespace spirv { /// Mapping from numeric MemRef memory spaces into SPIR-V symbolic ones. using MemorySpaceToStorageClassMap = DenseMap; +/// Returns the default map for targeting Vulkan-flavored SPIR-V. +MemorySpaceToStorageClassMap getDefaultVulkanStorageClassMap(); /// Type converter for converting numeric MemRef memory spaces into SPIR-V /// symbolic ones. diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td --- a/mlir/include/mlir/Conversion/Passes.td +++ b/mlir/include/mlir/Conversion/Passes.td @@ -543,9 +543,8 @@ let constructor = "mlir::createMapMemRefStorageClassPass()"; let dependentDialects = ["spirv::SPIRVDialect"]; let options = [ - Option<"mappings", "mappings", "std::string", /*default=*/"", - "A comma-separated list of memory space to storage class mappings; " - "for example, '0=StorageClass,1=Uniform,2=Workgroup'"> + Option<"clientAPI", "client-api", "std::string", /*default=*/"\"vulkan\"", + "The client API to use for populating mappings"> ]; } diff --git a/mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp b/mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp --- a/mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp +++ b/mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp @@ -16,6 +16,7 @@ #include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRVPass.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h" +#include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/Transforms/DialectConversion.h" #include "llvm/ADT/StringExtras.h" @@ -26,37 +27,47 @@ using namespace mlir; //===----------------------------------------------------------------------===// -// Utility Functions +// Mappings //===----------------------------------------------------------------------===// -/// Parses the memory space mapping string `memorySpaceMapStr `and writes the -/// mappings encoded inside to `memorySpaceMap`. -static bool -parseMappingStr(StringRef memorySpaceMapStr, - spirv::MemorySpaceToStorageClassMap &memorySpaceMap) { - memorySpaceMap.clear(); - SmallVector mappings; - llvm::SplitString(memorySpaceMapStr, mappings, ","); - for (StringRef mapping : mappings) { - StringRef key, value; - std::tie(key, value) = mapping.split('='); - unsigned space; - if (!llvm::to_integer(key, space)) { - LLVM_DEBUG(llvm::dbgs() - << "failed to parse mapping string key: " << key << "\n"); - memorySpaceMap.clear(); - return false; - } - Optional storage = spirv::symbolizeStorageClass(value); - if (!storage) { - LLVM_DEBUG(llvm::dbgs() - << "failed to parse mapping string value: " << value << "\n"); - memorySpaceMap.clear(); - return false; - } - memorySpaceMap[space] = *storage; - } - return true; +spirv::MemorySpaceToStorageClassMap spirv::getDefaultVulkanStorageClassMap() { +/// Mapping between SPIR-V storage classes to memref memory spaces. +/// +/// Note: memref does not have a defined semantics for each memory space; it +/// depends on the context where it is used. There are no particular reasons +/// behind the number assignments; we try to follow NVVM conventions and largely +/// give common storage classes a smaller number. +#define STORAGE_SPACE_MAP_LIST(MAP_FN) \ + MAP_FN(spirv::StorageClass::StorageBuffer, 0) \ + MAP_FN(spirv::StorageClass::Generic, 1) \ + MAP_FN(spirv::StorageClass::Workgroup, 3) \ + MAP_FN(spirv::StorageClass::Uniform, 4) \ + MAP_FN(spirv::StorageClass::Private, 5) \ + MAP_FN(spirv::StorageClass::Function, 6) \ + MAP_FN(spirv::StorageClass::PushConstant, 7) \ + MAP_FN(spirv::StorageClass::UniformConstant, 8) \ + MAP_FN(spirv::StorageClass::Input, 9) \ + MAP_FN(spirv::StorageClass::Output, 10) \ + MAP_FN(spirv::StorageClass::CrossWorkgroup, 11) \ + MAP_FN(spirv::StorageClass::AtomicCounter, 12) \ + MAP_FN(spirv::StorageClass::Image, 13) \ + MAP_FN(spirv::StorageClass::CallableDataKHR, 14) \ + MAP_FN(spirv::StorageClass::IncomingCallableDataKHR, 15) \ + MAP_FN(spirv::StorageClass::RayPayloadKHR, 16) \ + MAP_FN(spirv::StorageClass::HitAttributeKHR, 17) \ + MAP_FN(spirv::StorageClass::IncomingRayPayloadKHR, 18) \ + MAP_FN(spirv::StorageClass::ShaderRecordBufferKHR, 19) \ + MAP_FN(spirv::StorageClass::PhysicalStorageBuffer, 20) \ + MAP_FN(spirv::StorageClass::CodeSectionINTEL, 21) \ + MAP_FN(spirv::StorageClass::DeviceOnlyINTEL, 22) \ + MAP_FN(spirv::StorageClass::HostOnlyINTEL, 23) + +#define STORAGE_SPACE_MAP_FN(storage, space) {space, storage}, + + return {STORAGE_SPACE_MAP_LIST(STORAGE_SPACE_MAP_FN)}; + +#undef STORAGE_SPACE_MAP_FN +#undef STORAGE_SPACE_MAP_LIST } //===----------------------------------------------------------------------===// @@ -237,9 +248,11 @@ if (failed(Pass::initializeOptions(options))) return failure(); - if (!parseMappingStr(mappings, memorySpaceMap)) + if (clientAPI != "vulkan") return failure(); + memorySpaceMap = spirv::getDefaultVulkanStorageClassMap(); + LLVM_DEBUG({ llvm::dbgs() << "memory space to storage class mapping:\n"; if (memorySpaceMap.empty()) diff --git a/mlir/test/Conversion/MemRefToSPIRV/map-storage-class.mlir b/mlir/test/Conversion/MemRefToSPIRV/map-storage-class.mlir --- a/mlir/test/Conversion/MemRefToSPIRV/map-storage-class.mlir +++ b/mlir/test/Conversion/MemRefToSPIRV/map-storage-class.mlir @@ -1,74 +1,74 @@ -// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -map-memref-spirv-storage-class='mappings=0=StorageBuffer,1=Uniform,2=Workgroup,3=PushConstant' -verify-diagnostics %s -o - | FileCheck %s +// 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 -// Mappings: +// Vulkan Mappings: // 0 -> StorageBuffer (12) -// 2 -> Workgroup (4) -// 1 -> Uniform (2) -// 3 -> PushConstant (9) +// 1 -> Generic (8) +// 3 -> Workgroup (4) +// 4 -> Uniform (2) // TODO: create a StorageClass wrapper class so we can print the symbolc // storage class (instead of the backing IntegerAttr) and be able to // round trip the IR. -// CHECK-LABEL: func @operand_result +// VULKAN-LABEL: func @operand_result func.func @operand_result() { - // CHECK: memref + // VULKAN: memref %0 = "dialect.memref_producer"() : () -> (memref) - // CHECK: memref<4xi32, 2 : i32> + // VULKAN: memref<4xi32, 8 : i32> %1 = "dialect.memref_producer"() : () -> (memref<4xi32, 1>) - // CHECK: memref - %2 = "dialect.memref_producer"() : () -> (memref) - // CHECK: memref<*xf16, 9 : i32> - %3 = "dialect.memref_producer"() : () -> (memref<*xf16, 3>) + // VULKAN: memref + %2 = "dialect.memref_producer"() : () -> (memref) + // VULKAN: memref<*xf16, 2 : i32> + %3 = "dialect.memref_producer"() : () -> (memref<*xf16, 4>) "dialect.memref_consumer"(%0) : (memref) -> () - // CHECK: memref<4xi32, 2 : i32> + // VULKAN: memref<4xi32, 8 : i32> "dialect.memref_consumer"(%1) : (memref<4xi32, 1>) -> () - // CHECK: memref - "dialect.memref_consumer"(%2) : (memref) -> () - // CHECK: memref<*xf16, 9 : i32> - "dialect.memref_consumer"(%3) : (memref<*xf16, 3>) -> () + // VULKAN: memref + "dialect.memref_consumer"(%2) : (memref) -> () + // VULKAN: memref<*xf16, 2 : i32> + "dialect.memref_consumer"(%3) : (memref<*xf16, 4>) -> () return } // ----- -// CHECK-LABEL: func @type_attribute +// VULKAN-LABEL: func @type_attribute func.func @type_attribute() { - // CHECK: attr = memref + // VULKAN: attr = memref "dialect.memref_producer"() { attr = memref } : () -> () return } // ----- -// CHECK-LABEL: func @function_io +// VULKAN-LABEL: func @function_io func.func @function_io - // CHECK-SAME: (%{{.+}}: memref, %{{.+}}: memref<4xi32, 9 : i32>) - (%arg0: memref, %arg1: memref<4xi32, 3>) - // CHECK-SAME: -> (memref, memref<4xi32, 9 : i32>) - -> (memref, memref<4xi32, 3>) { - return %arg0, %arg1: memref, memref<4xi32, 3> + // VULKAN-SAME: (%{{.+}}: memref, %{{.+}}: memref<4xi32, 4 : i32>) + (%arg0: memref, %arg1: memref<4xi32, 3>) + // VULKAN-SAME: -> (memref, memref<4xi32, 4 : i32>) + -> (memref, memref<4xi32, 3>) { + return %arg0, %arg1: memref, memref<4xi32, 3> } // ----- -// CHECK: func @region +// VULKAN: func @region func.func @region(%cond: i1, %arg0: memref) { scf.if %cond { - // CHECK: "dialect.memref_consumer"(%{{.+}}) {attr = memref} - // CHECK-SAME: (memref) -> memref - %0 = "dialect.memref_consumer"(%arg0) { attr = memref } : (memref) -> (memref) + // VULKAN: "dialect.memref_consumer"(%{{.+}}) {attr = memref} + // VULKAN-SAME: (memref) -> memref + %0 = "dialect.memref_consumer"(%arg0) { attr = memref } : (memref) -> (memref) } return } // ----- -// CHECK-LABEL: func @non_memref_types +// VULKAN-LABEL: func @non_memref_types func.func @non_memref_types(%arg: f32) -> f32 { - // CHECK: "dialect.op"(%{{.+}}) {attr = 16 : i64} : (f32) -> f32 + // VULKAN: "dialect.op"(%{{.+}}) {attr = 16 : i64} : (f32) -> f32 %0 = "dialect.op"(%arg) { attr = 16 } : (f32) -> (f32) return %0 : f32 } @@ -77,6 +77,6 @@ func.func @missing_mapping() { // expected-error @+1 {{failed to legalize}} - %0 = "dialect.memref_producer"() : () -> (memref) + %0 = "dialect.memref_producer"() : () -> (memref) return }