Skip to content
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.

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

Use LLVMGPUTileandFuse instead of LLVMGPUVectorize for convolutions #19469

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft
Changes from all commits
Commits
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: 13 additions & 7 deletions compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1988,12 +1988,19 @@ static bool distributeToSquare(const int64_t oh, const int64_t ow,
// Convolution Pipeline Configuration
//====---------------------------------------------------------------------===//

static LogicalResult setConvolutionConfig(IREE::GPU::TargetAttr target,
linalg::LinalgOp linalgOp,
const int64_t bestTilingFactor) {
static LogicalResult setConvolutionConfig(
IREE::GPU::TargetAttr target, mlir::FunctionOpInterface entryPointFn,
linalg::LinalgOp linalgOp, const int64_t bestTilingFactor) {
if (!isa<linalg::Conv2DNhwcHwcfOp, linalg::Conv2DNchwFchwOp>(linalgOp)) {
return failure();
}
if (!clLLVMGPUVectorizePipeline) {
if (succeeded(IREE::GPU::setTileAndFuseLoweringConfig(target, entryPointFn,
linalgOp))) {
LDBG("Tile and fuse default config");
return success();
}
}
const bool isNCHW = isa<linalg::Conv2DNchwFchwOp>(*linalgOp);
const bool isNHWC = isa<linalg::Conv2DNhwcHwcfOp>(*linalgOp);

Expand Down Expand Up @@ -2076,9 +2083,8 @@ static LogicalResult setConvolutionConfig(IREE::GPU::TargetAttr target,
SmallVector<int64_t> windowTileSizes(4, 0);
windowTileSizes[ohIndex] = 1;
tileSizes.push_back(windowTileSizes);
auto funcOp = linalgOp->getParentOfType<mlir::FunctionOpInterface>();
return setOpConfigAndEntryPointFnTranslation(funcOp, linalgOp, tileSizes,
pipeline, workgroupSize);
return setOpConfigAndEntryPointFnTranslation(
entryPointFn, linalgOp, tileSizes, pipeline, workgroupSize);
}

//====---------------------------------------------------------------------===//
Expand Down Expand Up @@ -2132,7 +2138,7 @@ static LogicalResult setRootConfig(IREE::GPU::TargetAttr target,
LDBG("Warp Reduction Config");
return success();
}
if (succeeded(setConvolutionConfig(target, linalgOp, 16))) {
if (succeeded(setConvolutionConfig(target, entryPointFn, linalgOp, 16))) {
LDBG("Convolution Config");
return success();
}
Expand Down
Loading