Skip to content

[MLIR][GPU-LLVM] Add in-pass signature update option for opencl kernels #105664

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 18 commits into from
Oct 10, 2024
Merged

[MLIR][GPU-LLVM] Add in-pass signature update option for opencl kernels #105664

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
2e8465c
[MLIR][GPU-LLVM] Add in-pass signature update option for opencl kernels
kurapov-peter Aug 19, 2024
9238460
address review comment
kurapov-peter Aug 22, 2024
7212344
Update mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp
kurapov-peter Aug 26, 2024
1961dff
add debug type
kurapov-peter Aug 26, 2024
53baba2
Add func & memeref to llvm conversion and test a more complex caller/…
kurapov-peter Aug 29, 2024
4b18490
Update calling convention for func.func after the lowering
kurapov-peter Sep 11, 2024
dfbcd23
Update mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp
kurapov-peter Sep 11, 2024
b39e055
Rename OPENCL to OPENCL-CHECK
kurapov-peter Sep 13, 2024
d5519a5
remove calling convention updates to func.func
kurapov-peter Sep 19, 2024
598f9b1
Update mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp
kurapov-peter Sep 19, 2024
3fde4ae
Make forcing opencl address spaces the default behavior
kurapov-peter Sep 19, 2024
71006c6
put the comma back
kurapov-peter Sep 19, 2024
a1666d6
Insert integer attributes instead of the gpu address space
kurapov-peter Sep 19, 2024
c2b22c7
Avoid hardcoding default address space value
kurapov-peter Sep 20, 2024
a0a5a00
Fix 322-bit case
kurapov-peter Sep 20, 2024
abb6ef9
Drop unused includes
kurapov-peter Sep 23, 2024
0f79242
Merge branch 'main' into ocl-gpu-to-llvm-spv-signature-patching-3
kurapov-peter Sep 30, 2024
69ed255
Merge branch 'main' into ocl-gpu-to-llvm-spv-signature-patching-3
victor-eds Oct 9, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 49 additions & 0 deletions mlir/lib/Conversion/GPUToLLVMSPV/GPUToLLVMSPV.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/FormatVariadic.h"

#define DEBUG_TYPE "gpu-to-llvm-spv"

using namespace mlir;

namespace mlir {
Expand Down Expand Up @@ -316,6 +318,38 @@ struct GPUShuffleConversion final : ConvertOpToLLVMPattern<gpu::ShuffleOp> {
}
};

class MemorySpaceToOpenCLMemorySpaceConverter final : public TypeConverter {
public:
MemorySpaceToOpenCLMemorySpaceConverter(MLIRContext *ctx) {
addConversion([](Type t) { return t; });
addConversion([ctx](BaseMemRefType memRefType) -> std::optional<Type> {
// Attach global addr space attribute to memrefs with no addr space attr
Attribute memSpaceAttr = memRefType.getMemorySpace();
if (memSpaceAttr)
return std::nullopt;

unsigned globalAddrspace = storageClassToAddressSpace(
spirv::ClientAPI::OpenCL, spirv::StorageClass::CrossWorkgroup);
Attribute addrSpaceAttr =
IntegerAttr::get(IntegerType::get(ctx, 64), globalAddrspace);
if (auto rankedType = dyn_cast<MemRefType>(memRefType)) {
return MemRefType::get(memRefType.getShape(),
memRefType.getElementType(),
rankedType.getLayout(), addrSpaceAttr);
}
return UnrankedMemRefType::get(memRefType.getElementType(),
addrSpaceAttr);
});
addConversion([this](FunctionType type) {
auto inputs = llvm::map_to_vector(
type.getInputs(), [this](Type ty) { return convertType(ty); });
auto results = llvm::map_to_vector(
type.getResults(), [this](Type ty) { return convertType(ty); });
return FunctionType::get(type.getContext(), inputs, results);
});
}
};

//===----------------------------------------------------------------------===//
// Subgroup query ops.
//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -382,6 +416,21 @@ struct GPUToLLVMSPVConversionPass final
LLVMTypeConverter converter(context, options);
LLVMConversionTarget target(*context);

// Force OpenCL address spaces when they are not present
{
MemorySpaceToOpenCLMemorySpaceConverter converter(context);
AttrTypeReplacer replacer;
replacer.addReplacement([&converter](BaseMemRefType origType)
-> std::optional<BaseMemRefType> {
return converter.convertType<BaseMemRefType>(origType);
});

replacer.recursivelyReplaceElementsIn(getOperation(),
/*replaceAttrs=*/true,
/*replaceLocs=*/false,
/*replaceTypes=*/true);
}

target.addIllegalOp<gpu::BarrierOp, gpu::BlockDimOp, gpu::BlockIdOp,
gpu::GPUFuncOp, gpu::GlobalIdOp, gpu::GridDimOp,
gpu::LaneIdOp, gpu::NumSubgroupsOp, gpu::ReturnOp,
Expand Down
50 changes: 44 additions & 6 deletions mlir/test/Conversion/GPUToLLVMSPV/gpu-to-llvm-spv.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -444,20 +444,20 @@ gpu.module @kernels {
gpu.return
}

// CHECK-64: llvm.func spir_kernelcc @kernel_with_conv_args(%{{.*}}: i64, %{{.*}}: !llvm.ptr, %{{.*}}: !llvm.ptr, %{{.*}}: i64) attributes {gpu.kernel} {
// CHECK-32: llvm.func spir_kernelcc @kernel_with_conv_args(%{{.*}}: i32, %{{.*}}: !llvm.ptr, %{{.*}}: !llvm.ptr, %{{.*}}: i32) attributes {gpu.kernel} {
// CHECK-64: llvm.func spir_kernelcc @kernel_with_conv_args(%{{.*}}: i64, %{{.*}}: !llvm.ptr<1>, %{{.*}}: !llvm.ptr<1>, %{{.*}}: i64) attributes {gpu.kernel} {
// CHECK-32: llvm.func spir_kernelcc @kernel_with_conv_args(%{{.*}}: i32, %{{.*}}: !llvm.ptr<1>, %{{.*}}: !llvm.ptr<1>, %{{.*}}: i32) attributes {gpu.kernel} {
gpu.func @kernel_with_conv_args(%arg0: index, %arg1: memref<index>) kernel {
gpu.return
}

// CHECK-64: llvm.func spir_kernelcc @kernel_with_sized_memref(%{{.*}}: !llvm.ptr, %{{.*}}: !llvm.ptr, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: i64) attributes {gpu.kernel} {
// CHECK-32: llvm.func spir_kernelcc @kernel_with_sized_memref(%{{.*}}: !llvm.ptr, %{{.*}}: !llvm.ptr, %{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i32) attributes {gpu.kernel} {
// CHECK-64: llvm.func spir_kernelcc @kernel_with_sized_memref(%{{.*}}: !llvm.ptr<1>, %{{.*}}: !llvm.ptr<1>, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: i64) attributes {gpu.kernel} {
// CHECK-32: llvm.func spir_kernelcc @kernel_with_sized_memref(%{{.*}}: !llvm.ptr<1>, %{{.*}}: !llvm.ptr<1>, %{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i32) attributes {gpu.kernel} {
gpu.func @kernel_with_sized_memref(%arg0: memref<1xindex>) kernel {
gpu.return
}

// CHECK-64: llvm.func spir_kernelcc @kernel_with_ND_memref(%{{.*}}: !llvm.ptr, %{{.*}}: !llvm.ptr, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: i64) attributes {gpu.kernel} {
// CHECK-32: llvm.func spir_kernelcc @kernel_with_ND_memref(%{{.*}}: !llvm.ptr, %{{.*}}: !llvm.ptr, %{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i32) attributes {gpu.kernel} {
// CHECK-64: llvm.func spir_kernelcc @kernel_with_ND_memref(%{{.*}}: !llvm.ptr<1>, %{{.*}}: !llvm.ptr<1>, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: i64) attributes {gpu.kernel} {
// CHECK-32: llvm.func spir_kernelcc @kernel_with_ND_memref(%{{.*}}: !llvm.ptr<1>, %{{.*}}: !llvm.ptr<1>, %{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i32) attributes {gpu.kernel} {
gpu.func @kernel_with_ND_memref(%arg0: memref<128x128x128xindex>) kernel {
gpu.return
}
Expand Down Expand Up @@ -566,6 +566,44 @@ gpu.module @kernels {

// -----

gpu.module @kernels {
// CHECK: llvm.func spir_funccc @_Z12get_group_idj(i32)
// CHECK-LABEL: llvm.func spir_funccc @no_address_spaces(
// CHECK-SAME: %{{[a-zA-Z_][a-zA-Z0-9_]*}}: !llvm.ptr<1>
// CHECK-SAME: %{{[a-zA-Z_][a-zA-Z0-9_]*}}: !llvm.ptr<1>
// CHECK-SAME: %{{[a-zA-Z_][a-zA-Z0-9_]*}}: !llvm.ptr<1>
gpu.func @no_address_spaces(%arg0: memref<f32>, %arg1: memref<f32, #gpu.address_space<global>>, %arg2: memref<f32>) {
gpu.return
}

// CHECK-LABEL: llvm.func spir_kernelcc @no_address_spaces_complex(
// CHECK-SAME: %{{[a-zA-Z_][a-zA-Z0-9_]*}}: !llvm.ptr<1>
// CHECK-SAME: %{{[a-zA-Z_][a-zA-Z0-9_]*}}: !llvm.ptr<1>
// CHECK: func.call @no_address_spaces_callee(%{{[0-9]+}}, %{{[0-9]+}})
// CHECK-SAME: : (memref<2x2xf32, 1>, memref<4xf32, 1>)
gpu.func @no_address_spaces_complex(%arg0: memref<2x2xf32>, %arg1: memref<4xf32>) kernel {
func.call @no_address_spaces_callee(%arg0, %arg1) : (memref<2x2xf32>, memref<4xf32>) -> ()
gpu.return
}
// CHECK-LABEL: func.func @no_address_spaces_callee(
// CHECK-SAME: [[ARG0:%.*]]: memref<2x2xf32, 1>
// CHECK-SAME: [[ARG1:%.*]]: memref<4xf32, 1>
// CHECK: [[C0:%.*]] = llvm.mlir.constant(0 : i32) : i32
// CHECK: [[I0:%.*]] = llvm.call spir_funccc @_Z12get_group_idj([[C0]]) {
// CHECK-32: [[I1:%.*]] = builtin.unrealized_conversion_cast [[I0]] : i32 to index
// CHECK-64: [[I1:%.*]] = builtin.unrealized_conversion_cast [[I0]] : i64 to index
// CHECK: [[LD:%.*]] = memref.load [[ARG0]]{{\[}}[[I1]], [[I1]]{{\]}} : memref<2x2xf32, 1>
// CHECK: memref.store [[LD]], [[ARG1]]{{\[}}[[I1]]{{\]}} : memref<4xf32, 1>
func.func @no_address_spaces_callee(%arg0: memref<2x2xf32>, %arg1: memref<4xf32>) {
%block_id = gpu.block_id x
%0 = memref.load %arg0[%block_id, %block_id] : memref<2x2xf32>
memref.store %0, %arg1[%block_id] : memref<4xf32>
func.return
}
}

// -----

// Lowering of subgroup query operations

// CHECK-DAG: llvm.func spir_funccc @_Z18get_sub_group_size() -> i32 attributes {no_unwind, will_return}
Expand Down