Skip to content

Commit

Permalink
[NFC][GPU] Remove named op IGEMM lowering (#19885)
Browse files Browse the repository at this point in the history 
This was left behind when we switched to generic lowering in case the
broader support it brought caused issues down the pipeline but it does
not seem that is the case so we can remove the logic related to named op
lowering.

Signed-off-by: Nirvedh Meshram <nirvedh@gmail.com>
  • Loading branch information
@nirvedhmeshram
nirvedhmeshram authored Feb 3, 2025
1 parent aff23ae commit 539606c
Show file tree
Hide file tree
Showing 3 changed files with 11 additions and 126 deletions.
43 changes: 11 additions & 32 deletions compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,14 +31,6 @@

namespace mlir::iree_compiler::IREE::GPU {

// TODO (nirvedhmeshram) : This flag allows a lot more convolutions to use IGEMM
// so drop this flag after sufficient use with no issues.
llvm::cl::opt<bool> clGPUUseTileAndFuseGenericConvolution(
"iree-gpu-use-tile-and-fuse-generic-convolution",
llvm::cl::desc(
"enable the tile and fuse pipeline for generic convolutions"),
llvm::cl::init(true));

constexpr int64_t kCacheLineSizeBits = 128 * 8;
constexpr int64_t kPreferredCopyNumBits = 128;

Expand Down Expand Up @@ -383,35 +375,22 @@ setIGEMMConvolutionLoweringConfig(IREE::GPU::TargetAttr target,
return failure();

LDBG("IGEMM TileAndFuse Config");
FailureOr<SmallVector<AffineMap>> igemmContractionMaps;
FailureOr<SmallVector<int64_t>> igemmLoopBounds;
FailureOr<SmallVector<Value>> igemmOperands;
if (!clGPUUseTileAndFuseGenericConvolution) {
igemmContractionMaps = LinalgExt::getIGEMMContractionIndexingMaps(linalgOp);
igemmLoopBounds = LinalgExt::getIGEMMLoopBounds(linalgOp);
igemmOperands = LinalgExt::getIGEMMOperands(linalgOp);
} else {
FailureOr<LinalgExt::IGEMMGenericConvDetails> igemmGenericConvDetails =
LinalgExt::getIGEMMGenericConvDetails(linalgOp);
if (failed(igemmGenericConvDetails)) {
LDBG("Unsupported generic convolution type");
return failure();
}
igemmContractionMaps = igemmGenericConvDetails->igemmContractionMaps;
igemmLoopBounds = igemmGenericConvDetails->igemmLoopBounds;
igemmOperands = igemmGenericConvDetails->igemmOperands;
}

if (failed(igemmContractionMaps) || failed(igemmLoopBounds) ||
failed(igemmOperands)) {
FailureOr<LinalgExt::IGEMMGenericConvDetails> igemmGenericConvDetails =
LinalgExt::getIGEMMGenericConvDetails(linalgOp);
if (failed(igemmGenericConvDetails)) {
LDBG("Unsupported convolution type");
return failure();
}
SmallVector<AffineMap> igemmContractionMaps =
igemmGenericConvDetails->igemmContractionMaps;
SmallVector<int64_t> igemmLoopBounds =
igemmGenericConvDetails->igemmLoopBounds;
SmallVector<Value> igemmOperands = igemmGenericConvDetails->igemmOperands;

SmallVector<int64_t> bounds = igemmLoopBounds.value();
SmallVector<int64_t> bounds = igemmLoopBounds;
FailureOr<std::pair<LoweringConfigAttr, int64_t>> configAndWgSize =
getMatmulLoweringConfigAndWorkgroupSize(
bounds, igemmContractionMaps.value(), igemmOperands.value(), target);
getMatmulLoweringConfigAndWorkgroupSize(bounds, igemmContractionMaps,
igemmOperands, target);
if (failed(configAndWgSize)) {
return failure();
}
Expand Down
78 changes: 0 additions & 78 deletions compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -384,84 +384,6 @@ bool isGatherlikeOp(Operation *op) {
return hasTensorExtract;
}

FailureOr<SmallVector<AffineMap>>
getIGEMMContractionIndexingMaps(linalg::LinalgOp linalgOp) {
MLIRContext *ctx = linalgOp->getContext();
return llvm::TypeSwitch<Operation *, FailureOr<SmallVector<AffineMap>>>(
linalgOp.getOperation())
.Case<linalg::Conv2DNchwFchwOp>(
[&](linalg::Conv2DNchwFchwOp convOp) -> SmallVector<AffineMap> {
AffineExpr bDim, mDim, nDim0, nDim1, kDim;
bindDims(ctx, bDim, mDim, nDim0, nDim1, kDim);
auto lhsMap = AffineMap::get(5, 0, {mDim, kDim}, ctx);
auto rhsMap = AffineMap::get(5, 0, {bDim, nDim0, nDim1, kDim}, ctx);
auto resultMap =
AffineMap::get(5, 0, {bDim, mDim, nDim0, nDim1}, ctx);
return {lhsMap, rhsMap, resultMap};
})
.Case<linalg::Conv2DNhwcHwcfOp>(
[&](linalg::Conv2DNhwcHwcfOp convOp) -> SmallVector<AffineMap> {
AffineExpr bDim, m0Dim, m1Dim, nDim, kDim;
bindDims(ctx, bDim, m0Dim, m1Dim, nDim, kDim);
auto lhsMap = AffineMap::get(5, 0, {bDim, m0Dim, m1Dim, kDim}, ctx);
auto rhsMap = AffineMap::get(5, 0, {kDim, nDim}, ctx);
auto resultMap =
AffineMap::get(5, 0, {bDim, m0Dim, m1Dim, nDim}, ctx);
return {lhsMap, rhsMap, resultMap};
})
.Default([](Operation *) { return failure(); });
}

FailureOr<SmallVector<int64_t>> getIGEMMLoopBounds(linalg::LinalgOp linalgOp) {
return llvm::TypeSwitch<Operation *, FailureOr<SmallVector<int64_t>>>(
linalgOp.getOperation())
.Case<linalg::Conv2DNchwFchwOp>(
[&](linalg::Conv2DNchwFchwOp convOp) -> SmallVector<int64_t> {
auto filterType =
cast<RankedTensorType>(convOp.getOperandTypes()[1]);
auto accType = cast<RankedTensorType>(convOp.getResultTypes()[0]);
const int64_t B = accType.getDimSize(0);
const int64_t N0 = accType.getDimSize(2);
const int64_t N1 = accType.getDimSize(3);
const int64_t M = filterType.getDimSize(0);
const int64_t K = filterType.getDimSize(1) *
filterType.getDimSize(2) *
filterType.getDimSize(3);
return {B, M, N0, N1, K};
})
.Case<linalg::Conv2DNhwcHwcfOp>(
[&](linalg::Conv2DNhwcHwcfOp convOp) -> SmallVector<int64_t> {
auto filterType =
cast<RankedTensorType>(convOp.getOperandTypes()[1]);
auto accType = cast<RankedTensorType>(convOp.getResultTypes()[0]);
const int64_t B = accType.getDimSize(0);
const int64_t M0 = accType.getDimSize(1);
const int64_t M1 = accType.getDimSize(2);
const int64_t N = accType.getDimSize(3);
const int64_t K = filterType.getDimSize(0) *
filterType.getDimSize(1) *
filterType.getDimSize(2);
return {B, M0, M1, N, K};
})
.Default([](Operation *) { return failure(); });
}

FailureOr<SmallVector<Value>> getIGEMMOperands(linalg::LinalgOp linalgOp) {
return llvm::TypeSwitch<Operation *, FailureOr<SmallVector<Value>>>(
linalgOp.getOperation())
.Case<linalg::Conv2DNchwFchwOp>(
[&](linalg::Conv2DNchwFchwOp convOp) -> SmallVector<Value> {
return {convOp.getOperands()[1], convOp.getOperands()[0],
convOp.getOperands()[2]};
})
.Case<linalg::Conv2DNhwcHwcfOp>(
[&](linalg::Conv2DNhwcHwcfOp convOp) -> SmallVector<Value> {
return {convOp.getOperands()[0], convOp.getOperands()[1],
convOp.getOperands()[2]};
})
.Default([](Operation *) { return failure(); });
}

FailureOr<IGEMMGenericConvDetails>
getIGEMMGenericConvDetails(linalg::LinalgOp linalgOp) {

Expand Down
16 changes: 0 additions & 16 deletions compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -145,22 +145,6 @@ SmallVector<AffineMap> convertDimsToSymbols(MLIRContext *context,
unsigned numDims,
unsigned numSymbols);

/// Returns the indexing maps array for a convolution operation with IGEMM
/// indexing. The resulting indexing maps should represent the indexing of some
/// contraction that computes the equivalent IGEMM matmul of the convolution.
FailureOr<SmallVector<AffineMap>>
getIGEMMContractionIndexingMaps(linalg::LinalgOp linalgOp);

/// Returns the loop bounds of a convolution op with IGEMM indexing. This
/// function assumes the same ordering of dimensions as
/// getIGEMMContractionIndexingMaps;
FailureOr<SmallVector<int64_t>> getIGEMMLoopBounds(linalg::LinalgOp linalgOp);

/// Returns the operand list for a convolution with IGEMM indexing. This is
/// used to determine which inputs are the lhs and rhs, since depending on the
/// layout, the order can be different (e.g., NCHW has the lhs and rhs swapped).
FailureOr<SmallVector<Value>> getIGEMMOperands(linalg::LinalgOp linalgOp);

/// Struct that holds inferred IGEMM details for a convolution operation.
struct IGEMMGenericConvDetails {
/// The indexing maps array for a convolution operation with IGEMM
Expand Down

0 comments on commit 539606c

Please sign in to comment.