Skip to content

Commit 0718e3a

Browse files
author
Denis Khalikov
committed
[mlir][vulkan-runner] Add support for 3D memrefs.
Summary: Add support for 3D memrefs in mlir-vulkan-runner and simple test. Differential Revision: https://reviews.llvm.org/D77157
1 parent 6aecf0c commit 0718e3a

File tree

4 files changed

+97
-16
lines changed

4 files changed

+97
-16
lines changed

‎mlir/lib/Conversion/GPUToVulkan/ConvertGPULaunchFuncToVulkanLaunchFunc.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ class ConvertGpuLaunchFuncToVulkanLaunchFunc
6060
// TODO(denis0x0D): Handle other types.
6161
if (auto memRefType = type.dyn_cast_or_null<MemRefType>())
6262
return memRefType.hasRank() &&
63-
(memRefType.getRank() == 1 || memRefType.getRank() == 2);
63+
(memRefType.getRank() >= 1 && memRefType.getRank() <= 3);
6464
return false;
6565
}
6666

‎mlir/lib/Conversion/GPUToVulkan/ConvertLaunchFuncToVulkanCalls.cpp‎

Lines changed: 23 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ using namespace mlir;
3030

3131
static constexpr const char *kBindMemRef1DFloat = "bindMemRef1DFloat";
3232
static constexpr const char *kBindMemRef2DFloat = "bindMemRef2DFloat";
33+
static constexpr const char *kBindMemRef3DFloat = "bindMemRef3DFloat";
3334
static constexpr const char *kCInterfaceVulkanLaunch =
3435
"_mlir_ciface_vulkanLaunch";
3536
static constexpr const char *kDeinitVulkan = "deinitVulkan";
@@ -76,10 +77,12 @@ class VulkanLaunchFuncToVulkanCallsPass
7677
llvmPointerType = LLVM::LLVMType::getInt8PtrTy(llvmDialect);
7778
llvmInt32Type = LLVM::LLVMType::getInt32Ty(llvmDialect);
7879
llvmInt64Type = LLVM::LLVMType::getInt64Ty(llvmDialect);
79-
initializeMemRefTypes();
80+
llvmMemRef1DFloat = getMemRefType(1);
81+
llvmMemRef2DFloat = getMemRefType(2);
82+
llvmMemRef3DFloat = getMemRefType(3);
8083
}
8184

82-
void initializeMemRefTypes() {
85+
LLVM::LLVMType getMemRefType(uint32_t rank) {
8386
// According to the MLIR doc memref argument is converted into a
8487
// pointer-to-struct argument of type:
8588
// template <typename Elem, size_t Rank>
@@ -91,22 +94,15 @@ class VulkanLaunchFuncToVulkanCallsPass
9194
// int64_t strides[Rank]; // omitted when rank == 0
9295
// };
9396
auto llvmPtrToFloatType = getFloatType().getPointerTo();
94-
auto llvmArrayOneElementSizeType =
95-
LLVM::LLVMType::getArrayTy(getInt64Type(), 1);
96-
auto llvmArrayTwoElementSizeType =
97-
LLVM::LLVMType::getArrayTy(getInt64Type(), 2);
97+
auto llvmArrayRankElementSizeType =
98+
LLVM::LLVMType::getArrayTy(getInt64Type(), rank);
9899

99-
// Create a type `!llvm<"{ float*, float*, i64, [1 x i64], [1 x i64]}">`.
100-
llvmMemRef1DFloat = LLVM::LLVMType::getStructTy(
100+
// Create a type
101+
// `!llvm<"{ float*, float*, i64, [`rank` x i64], [`rank` x i64]}">`.
102+
return LLVM::LLVMType::getStructTy(
101103
llvmDialect,
102104
{llvmPtrToFloatType, llvmPtrToFloatType, getInt64Type(),
103-
llvmArrayOneElementSizeType, llvmArrayOneElementSizeType});
104-
105-
// Create a type `!llvm<"{ float*, float*, i64, [2 x i64], [2 x i64]}">`.
106-
llvmMemRef2DFloat = LLVM::LLVMType::getStructTy(
107-
llvmDialect,
108-
{llvmPtrToFloatType, llvmPtrToFloatType, getInt64Type(),
109-
llvmArrayTwoElementSizeType, llvmArrayTwoElementSizeType});
105+
llvmArrayRankElementSizeType, llvmArrayRankElementSizeType});
110106
}
111107

112108
LLVM::LLVMType getFloatType() { return llvmFloatType; }
@@ -116,6 +112,7 @@ class VulkanLaunchFuncToVulkanCallsPass
116112
LLVM::LLVMType getInt64Type() { return llvmInt64Type; }
117113
LLVM::LLVMType getMemRef1DFloat() { return llvmMemRef1DFloat; }
118114
LLVM::LLVMType getMemRef2DFloat() { return llvmMemRef2DFloat; }
115+
LLVM::LLVMType getMemRef3DFloat() { return llvmMemRef3DFloat; }
119116

120117
/// Creates a LLVM global for the given `name`.
121118
Value createEntryPointNameConstant(StringRef name, Location loc,
@@ -164,6 +161,7 @@ class VulkanLaunchFuncToVulkanCallsPass
164161
LLVM::LLVMType llvmInt64Type;
165162
LLVM::LLVMType llvmMemRef1DFloat;
166163
LLVM::LLVMType llvmMemRef2DFloat;
164+
LLVM::LLVMType llvmMemRef3DFloat;
167165

168166
// TODO: Use an associative array to support multiple vulkan launch calls.
169167
std::pair<StringAttr, StringAttr> spirvAttributes;
@@ -335,6 +333,16 @@ void VulkanLaunchFuncToVulkanCallsPass::declareVulkanFunctions(Location loc) {
335333
/*isVarArg=*/false));
336334
}
337335

336+
if (!module.lookupSymbol(kBindMemRef3DFloat)) {
337+
builder.create<LLVM::LLVMFuncOp>(
338+
loc, kBindMemRef3DFloat,
339+
LLVM::LLVMType::getFunctionTy(getVoidType(),
340+
{getPointerType(), getInt32Type(),
341+
getInt32Type(),
342+
getMemRef3DFloat().getPointerTo()},
343+
/*isVarArg=*/false));
344+
}
345+
338346
if (!module.lookupSymbol(kInitVulkan)) {
339347
builder.create<LLVM::LLVMFuncOp>(
340348
loc, kInitVulkan,

‎mlir/test/mlir-vulkan-runner/subf.mlir‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
// RUN: mlir-vulkan-runner %s --shared-libs=%vulkan_wrapper_library_dir/libvulkan-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
2+
3+
// CHECK-COUNT-32: [2.2, 2.2, 2.2, 2.2]
4+
module attributes {
5+
gpu.container_module,
6+
spv.target_env = #spv.target_env<
7+
#spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
8+
{max_compute_workgroup_invocations = 128 : i32,
9+
max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
10+
} {
11+
gpu.module @kernels {
12+
gpu.func @kernel_sub(%arg0 : memref<8x4x4xf32>, %arg1 : memref<4x4xf32>, %arg2 : memref<8x4x4xf32>)
13+
attributes {gpu.kernel, spv.entry_point_abi = {local_size = dense<[1, 1, 1]>: vector<3xi32>}} {
14+
%x = "gpu.block_id"() {dimension = "x"} : () -> index
15+
%y = "gpu.block_id"() {dimension = "y"} : () -> index
16+
%z = "gpu.block_id"() {dimension = "z"} : () -> index
17+
%1 = load %arg0[%x, %y, %z] : memref<8x4x4xf32>
18+
%2 = load %arg1[%y, %z] : memref<4x4xf32>
19+
%3 = subf %1, %2 : f32
20+
store %3, %arg2[%x, %y, %z] : memref<8x4x4xf32>
21+
gpu.return
22+
}
23+
}
24+
25+
func @main() {
26+
%arg0 = alloc() : memref<8x4x4xf32>
27+
%arg1 = alloc() : memref<4x4xf32>
28+
%arg2 = alloc() : memref<8x4x4xf32>
29+
%0 = constant 0 : i32
30+
%1 = constant 1 : i32
31+
%2 = constant 2 : i32
32+
%value0 = constant 0.0 : f32
33+
%value1 = constant 3.3 : f32
34+
%value2 = constant 1.1 : f32
35+
%arg3 = memref_cast %arg0 : memref<8x4x4xf32> to memref<?x?x?xf32>
36+
%arg4 = memref_cast %arg1 : memref<4x4xf32> to memref<?x?xf32>
37+
%arg5 = memref_cast %arg2 : memref<8x4x4xf32> to memref<?x?x?xf32>
38+
call @fillResource3DFloat(%arg3, %value1) : (memref<?x?x?xf32>, f32) -> ()
39+
call @fillResource2DFloat(%arg4, %value2) : (memref<?x?xf32>, f32) -> ()
40+
call @fillResource3DFloat(%arg5, %value0) : (memref<?x?x?xf32>, f32) -> ()
41+
42+
%cst1 = constant 1 : index
43+
%cst4 = constant 4 : index
44+
%cst8 = constant 8 : index
45+
"gpu.launch_func"(%cst8, %cst4, %cst4, %cst1, %cst1, %cst1, %arg0, %arg1, %arg2) { kernel = "kernel_sub", kernel_module = @kernels }
46+
: (index, index, index, index, index, index, memref<8x4x4xf32>, memref<4x4xf32>, memref<8x4x4xf32>) -> ()
47+
%arg6 = memref_cast %arg5 : memref<?x?x?xf32> to memref<*xf32>
48+
call @print_memref_f32(%arg6) : (memref<*xf32>) -> ()
49+
return
50+
}
51+
func @fillResource2DFloat(%0 : memref<?x?xf32>, %1 : f32)
52+
func @fillResource3DFloat(%0 : memref<?x?x?xf32>, %1 : f32)
53+
func @print_memref_f32(%ptr : memref<*xf32>)
54+
}

‎mlir/tools/mlir-vulkan-runner/vulkan-runtime-wrappers.cpp‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,18 @@ void bindMemRef2DFloat(void *vkRuntimeManager, DescriptorSetIndex setIndex,
123123
->setResourceData(setIndex, bindIndex, memBuffer);
124124
}
125125

126+
/// Binds the given 3D float memref to the given descriptor set and descriptor
127+
/// index.
128+
void bindMemRef3DFloat(void *vkRuntimeManager, DescriptorSetIndex setIndex,
129+
BindingIndex bindIndex,
130+
MemRefDescriptor<float, 3> *ptr) {
131+
VulkanHostMemoryBuffer memBuffer{
132+
ptr->allocated, static_cast<uint32_t>(ptr->sizes[0] * ptr->sizes[1] *
133+
ptr->sizes[2] * sizeof(float))};
134+
reinterpret_cast<VulkanRuntimeManager *>(vkRuntimeManager)
135+
->setResourceData(setIndex, bindIndex, memBuffer);
136+
}
137+
126138
/// Fills the given 1D float memref with the given float value.
127139
void _mlir_ciface_fillResource1DFloat(MemRefDescriptor<float, 1> *ptr, // NOLINT
128140
float value) {
@@ -134,4 +146,11 @@ void _mlir_ciface_fillResource2DFloat(MemRefDescriptor<float, 2> *ptr, // NOLINT
134146
float value) {
135147
std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1], value);
136148
}
149+
150+
/// Fills the given 3D float memref with the given float value.
151+
void _mlir_ciface_fillResource3DFloat(MemRefDescriptor<float, 3> *ptr, // NOLINT
152+
float value) {
153+
std::fill_n(ptr->allocated, ptr->sizes[0] * ptr->sizes[1] * ptr->sizes[2],
154+
value);
155+
}
137156
}

0 commit comments

Comments
 (0)