diff options
author | Lei Zhang <antiagainst@google.com> | 2022-08-05 12:24:14 -0400 |
---|---|---|
committer | Lei Zhang <antiagainst@google.com> | 2022-08-05 12:30:14 -0400 |
commit | 1f7544a6797a55a230b836141feb7e610f3535d9 (patch) | |
tree | 98ca04232009349cd2a63b49be4d5d0be21357a5 | |
parent | 713f85d5952ab27d474aba2a960a893b7e7e438d (diff) |
[mlir][spirv] Add default Vulkan memory space to storage class mappinglinaro-local/ci/tcwg_kernel/llvm-master-arm-mainline-allnoconfig
Reviewed By: ThomasRaoux, kuhar
Differential Revision: https://reviews.llvm.org/D131128
4 files changed, 79 insertions, 65 deletions
diff --git a/mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h b/mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h index e7b4b7ab23de..60730246a7c6 100644 --- a/mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h +++ b/mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h @@ -23,6 +23,8 @@ class SPIRVTypeConverter; namespace spirv { /// Mapping from numeric MemRef memory spaces into SPIR-V symbolic ones. using MemorySpaceToStorageClassMap = DenseMap<unsigned, spirv::StorageClass>; +/// 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 index ed5348109484..88010d4d029c 100644 --- a/mlir/include/mlir/Conversion/Passes.td +++ b/mlir/include/mlir/Conversion/Passes.td @@ -543,9 +543,8 @@ def MapMemRefStorageClass : Pass<"map-memref-spirv-storage-class", "ModuleOp"> { 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 index ba0003ce5083..2d8504093e3c 100644 --- 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<StringRef> 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<spirv::StorageClass> 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 @@ LogicalResult MapMemRefStorageClassPass::initializeOptions(StringRef options) { 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 index 5a2fe262110d..5f474208ec47 100644 --- 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<f32, 12 : i32> + // VULKAN: memref<f32, 12 : i32> %0 = "dialect.memref_producer"() : () -> (memref<f32>) - // CHECK: memref<4xi32, 2 : i32> + // VULKAN: memref<4xi32, 8 : i32> %1 = "dialect.memref_producer"() : () -> (memref<4xi32, 1>) - // CHECK: memref<?x4xf16, 4 : i32> - %2 = "dialect.memref_producer"() : () -> (memref<?x4xf16, 2>) - // CHECK: memref<*xf16, 9 : i32> - %3 = "dialect.memref_producer"() : () -> (memref<*xf16, 3>) + // VULKAN: memref<?x4xf16, 4 : i32> + %2 = "dialect.memref_producer"() : () -> (memref<?x4xf16, 3>) + // VULKAN: memref<*xf16, 2 : i32> + %3 = "dialect.memref_producer"() : () -> (memref<*xf16, 4>) "dialect.memref_consumer"(%0) : (memref<f32>) -> () - // CHECK: memref<4xi32, 2 : i32> + // VULKAN: memref<4xi32, 8 : i32> "dialect.memref_consumer"(%1) : (memref<4xi32, 1>) -> () - // CHECK: memref<?x4xf16, 4 : i32> - "dialect.memref_consumer"(%2) : (memref<?x4xf16, 2>) -> () - // CHECK: memref<*xf16, 9 : i32> - "dialect.memref_consumer"(%3) : (memref<*xf16, 3>) -> () + // VULKAN: memref<?x4xf16, 4 : i32> + "dialect.memref_consumer"(%2) : (memref<?x4xf16, 3>) -> () + // 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<i32, 2 : i32> + // VULKAN: attr = memref<i32, 8 : i32> "dialect.memref_producer"() { attr = memref<i32, 1> } : () -> () return } // ----- -// CHECK-LABEL: func @function_io +// VULKAN-LABEL: func @function_io func.func @function_io - // CHECK-SAME: (%{{.+}}: memref<f64, 4 : i32>, %{{.+}}: memref<4xi32, 9 : i32>) - (%arg0: memref<f64, 2>, %arg1: memref<4xi32, 3>) - // CHECK-SAME: -> (memref<f64, 4 : i32>, memref<4xi32, 9 : i32>) - -> (memref<f64, 2>, memref<4xi32, 3>) { - return %arg0, %arg1: memref<f64, 2>, memref<4xi32, 3> + // VULKAN-SAME: (%{{.+}}: memref<f64, 8 : i32>, %{{.+}}: memref<4xi32, 4 : i32>) + (%arg0: memref<f64, 1>, %arg1: memref<4xi32, 3>) + // VULKAN-SAME: -> (memref<f64, 8 : i32>, memref<4xi32, 4 : i32>) + -> (memref<f64, 1>, memref<4xi32, 3>) { + return %arg0, %arg1: memref<f64, 1>, memref<4xi32, 3> } // ----- -// CHECK: func @region +// VULKAN: func @region func.func @region(%cond: i1, %arg0: memref<f32, 1>) { scf.if %cond { - // CHECK: "dialect.memref_consumer"(%{{.+}}) {attr = memref<i64, 4 : i32>} - // CHECK-SAME: (memref<f32, 2 : i32>) -> memref<f32, 2 : i32> - %0 = "dialect.memref_consumer"(%arg0) { attr = memref<i64, 2> } : (memref<f32, 1>) -> (memref<f32, 1>) + // VULKAN: "dialect.memref_consumer"(%{{.+}}) {attr = memref<i64, 4 : i32>} + // VULKAN-SAME: (memref<f32, 8 : i32>) -> memref<f32, 8 : i32> + %0 = "dialect.memref_consumer"(%arg0) { attr = memref<i64, 3> } : (memref<f32, 1>) -> (memref<f32, 1>) } 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 @non_memref_types(%arg: f32) -> f32 { func.func @missing_mapping() { // expected-error @+1 {{failed to legalize}} - %0 = "dialect.memref_producer"() : () -> (memref<f32, 128>) + %0 = "dialect.memref_producer"() : () -> (memref<f32, 2>) return } |