Skip to content

[mlir][gpu] Use known_block_size to set maxntid for NVVM target #77301

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 1 commit into from
Jan 8, 2024
Merged

[mlir][gpu] Use known_block_size to set maxntid for NVVM target #77301

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
20 changes: 19 additions & 1 deletion mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,8 +85,26 @@ GPUFuncOpLowering::matchAndRewrite(gpu::GPUFuncOp gpuFuncOp, OpAdaptor adaptor,
// Add a dialect specific kernel attribute in addition to GPU kernel
// attribute. The former is necessary for further translation while the
// latter is expected by gpu.launch_func.
if (gpuFuncOp.isKernel())
if (gpuFuncOp.isKernel()) {
attributes.emplace_back(kernelAttributeName, rewriter.getUnitAttr());

// Set the block size attribute if it is present.
if (kernelBlockSizeAttributeName.has_value()) {
std::optional<int32_t> dimX =
gpuFuncOp.getKnownBlockSize(gpu::Dimension::x);
std::optional<int32_t> dimY =
gpuFuncOp.getKnownBlockSize(gpu::Dimension::y);
std::optional<int32_t> dimZ =
gpuFuncOp.getKnownBlockSize(gpu::Dimension::z);
if (dimX.has_value() || dimY.has_value() || dimZ.has_value()) {
// If any of the dimensions are missing, fill them in with 1.
attributes.emplace_back(
kernelBlockSizeAttributeName.value(),
rewriter.getI32ArrayAttr(
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we use a DenseArray?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure, I can use DenseArray and I also think I should.

Some adjustments are needed here, and need to change some test. I can do it in a separate PR if it's okay

{dimX.value_or(1), dimY.value_or(1), dimZ.value_or(1)}));
}
}
}
auto llvmFuncOp = rewriter.create<LLVM::LLVMFuncOp>(
gpuFuncOp.getLoc(), gpuFuncOp.getName(), funcType,
LLVM::Linkage::External, /*dsoLocal=*/false, /*cconv=*/LLVM::CConv::C,
Expand Down
13 changes: 9 additions & 4 deletions mlir/lib/Conversion/GPUCommon/GPUOpsLowering.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,13 +36,15 @@ struct GPUDynamicSharedMemoryOpLowering
};

struct GPUFuncOpLowering : ConvertOpToLLVMPattern<gpu::GPUFuncOp> {
GPUFuncOpLowering(const LLVMTypeConverter &converter,
unsigned allocaAddrSpace, unsigned workgroupAddrSpace,
StringAttr kernelAttributeName)
GPUFuncOpLowering(
const LLVMTypeConverter &converter, unsigned allocaAddrSpace,
unsigned workgroupAddrSpace, StringAttr kernelAttributeName,
std::optional<StringAttr> kernelBlockSizeAttributeName = std::nullopt)
: ConvertOpToLLVMPattern<gpu::GPUFuncOp>(converter),
allocaAddrSpace(allocaAddrSpace),
workgroupAddrSpace(workgroupAddrSpace),
kernelAttributeName(kernelAttributeName) {}
kernelAttributeName(kernelAttributeName),
kernelBlockSizeAttributeName(kernelBlockSizeAttributeName) {}

LogicalResult
matchAndRewrite(gpu::GPUFuncOp gpuFuncOp, OpAdaptor adaptor,
Expand All @@ -56,6 +58,9 @@ struct GPUFuncOpLowering : ConvertOpToLLVMPattern<gpu::GPUFuncOp> {

/// The attribute name to use instead of `gpu.kernel`.
StringAttr kernelAttributeName;

/// The attribute name to to set block size
std::optional<StringAttr> kernelBlockSizeAttributeName;
};

/// The lowering of gpu.printf to a call to HIP hostcalls
Expand Down
4 changes: 3 additions & 1 deletion mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -352,7 +352,9 @@ void mlir::populateGpuToNVVMConversionPatterns(LLVMTypeConverter &converter,
/*workgroupAddrSpace=*/
static_cast<unsigned>(NVVM::NVVMMemorySpace::kSharedMemorySpace),
StringAttr::get(&converter.getContext(),
NVVM::NVVMDialect::getKernelFuncAttrName()));
NVVM::NVVMDialect::getKernelFuncAttrName()),
StringAttr::get(&converter.getContext(),
NVVM::NVVMDialect::getMaxntidAttrName()));

populateOpPatterns<math::AbsFOp>(converter, patterns, "__nv_fabsf",
"__nv_fabs");
Expand Down
9 changes: 9 additions & 0 deletions mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -627,6 +627,15 @@ gpu.module @test_module_31 {
}
}

gpu.module @gpumodule {
// CHECK-LABEL: func @kernel_with_block_size()
// CHECK: attributes {gpu.kernel, gpu.known_block_size = array<i32: 128, 1, 1>, nvvm.kernel, nvvm.maxntid = [128 : i32, 1 : i32, 1 : i32]}
gpu.func @kernel_with_block_size() kernel attributes {gpu.known_block_size = array<i32: 128, 1, 1>} {
gpu.return
}
}


module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%toplevel_module: !transform.any_op {transform.readonly}) {
%gpu_module = transform.structured.match ops{["gpu.module"]} in %toplevel_module
Expand Down