Skip to content

Commit acbd6de

Browse files
hanhanWhanhanW
authored
[DT][NFCI] Implement getOffsetsSizesStrides for GPU padding resolver. (#22339)
The revision implements the interface method with the existing logic; it allows us to remove padding specific load/store patterns, because they use interfaces. No new tests are added since it aims to be a refactor PR and tests should be covered by [materialize_encoding_into_padding.mlir](https://github.com/iree-org/iree/blob/3b1730e5fa15a520505000b3f3232bd2a230793b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_into_padding.mlir) It is a step towards #20160 --------- Signed-off-by: hanhanW <[email protected]>
1 parent 3d3e9c6 commit acbd6de

File tree

4 files changed

+31
-122
lines changed

4 files changed

+31
-122
lines changed

‎compiler/src/iree/compiler/Codegen/Common/EncodingUtils.cpp‎

Lines changed: 3 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -5,19 +5,11 @@
55
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
66

77
#include "iree/compiler/Codegen/Common/EncodingUtils.h"
8-
#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h"
9-
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
108
#include "iree/compiler/Codegen/Utils/EncodingUtils.h"
11-
#include "iree/compiler/Codegen/Utils/Utils.h"
129
#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
1310
#include "iree/compiler/Dialect/TensorExt/IR/TensorExtTypes.h"
14-
#include "llvm/Support/Debug.h"
15-
#include "mlir/Dialect/Tensor/IR/Tensor.h"
16-
#include "mlir/Dialect/Utils/IndexingUtils.h"
1711
#include "mlir/IR/BuiltinAttributes.h"
1812

19-
#include <optional>
20-
2113
#define DEBUG_TYPE "iree-codegen-encoding-utils"
2214

2315
namespace mlir::iree_compiler {
@@ -124,24 +116,9 @@ LogicalResult MaterializeEncodingTypeConverter::getOffsetsSizesStrides(
124116
SmallVectorImpl<OpFoldResult> &newOffsets,
125117
SmallVectorImpl<OpFoldResult> &newSizes,
126118
SmallVectorImpl<OpFoldResult> &newStrides) const {
127-
auto boundType = dyn_cast<RankedTensorType>(type.getBoundType());
128-
if (!boundType || !boundType.getEncoding()) {
129-
return failure();
130-
}
131-
// TODO(jornt): The isa<IREE::GPU::GPUPaddingResolverAttr> check is
132-
// needed because PaddingAttr is a serializable attribute, but it
133-
// relies on its own type conversion for now. Once GPUPaddingResolverAttr
134-
// implements `getOffsetsSizesStrides`, this can be removed.
135-
if (!isa<IREE::GPU::GPUPaddingResolverAttr>(getLayoutAttr())) {
136-
return getLayoutAttr().getOffsetsSizesStrides(
137-
builder, loc, type, dynamicDims, offsets, sizes, strides, newOffsets,
138-
newSizes, newStrides);
139-
}
140-
auto boundTensorType = cast<RankedTensorType>(type.getBoundType());
141-
newSizes = getMixedValues(boundTensorType.getShape(), dynamicDims, builder);
142-
newOffsets.resize(newSizes.size(), builder.getIndexAttr(0));
143-
newStrides.resize(newSizes.size(), builder.getIndexAttr(1));
144-
return success();
119+
return getLayoutAttr().getOffsetsSizesStrides(
120+
builder, loc, type, dynamicDims, offsets, sizes, strides, newOffsets,
121+
newSizes, newStrides);
145122
}
146123

147124
} // namespace mlir::iree_compiler

‎compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoPadding.cpp‎

Lines changed: 1 addition & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -116,95 +116,6 @@ struct MaterializePadEncodingTypeConverter final
116116
}
117117
};
118118

119-
/// Pattern to convert `iree_tensor_ext.dispatch.tensor.load` operation when
120-
/// materializing the encoding.
121-
struct MaterializeFlowDispatchTensorLoadOp final
122-
: OpConversionPattern<IREE::TensorExt::DispatchTensorLoadOp> {
123-
using Base::Base;
124-
125-
LogicalResult
126-
matchAndRewrite(IREE::TensorExt::DispatchTensorLoadOp loadOp,
127-
OpAdaptor adaptor,
128-
ConversionPatternRewriter &rewriter) const override {
129-
// Only handle operations where the load covers the entire
130-
// `!iree_tensor_ext.dispatch.tensor` type.
131-
if (!loadOp.isLoadOfWholeSource()) {
132-
return rewriter.notifyMatchFailure(loadOp, "unhandled partial loads");
133-
}
134-
135-
auto &typeConverter =
136-
*getTypeConverter<MaterializePadEncodingTypeConverter>();
137-
IREE::TensorExt::DispatchTensorType sourceType = loadOp.getSourceType();
138-
auto boundTensorType = cast<RankedTensorType>(sourceType.getBoundType());
139-
if (!typeConverter.hasNonZeroPadding(boundTensorType)) {
140-
// Let the Nop pattern handle this.
141-
return rewriter.notifyMatchFailure(loadOp, "no padding applied");
142-
}
143-
144-
auto paddedType =
145-
typeConverter.convertType<RankedTensorType>(boundTensorType);
146-
assert(paddedType != boundTensorType && "Expected conversion with padding");
147-
148-
SmallVector<OpFoldResult> newMixedSizes =
149-
getMixedValues(paddedType.getShape(), loadOp.getSourceDims(), rewriter);
150-
151-
SmallVector<OpFoldResult> newOffsets(newMixedSizes.size(),
152-
rewriter.getIndexAttr(0));
153-
SmallVector<OpFoldResult> newStrides(newMixedSizes.size(),
154-
rewriter.getIndexAttr(1));
155-
SmallVector<OpFoldResult> extractSizes = getMixedValues(
156-
boundTensorType.getShape(), loadOp.getSourceDims(), rewriter);
157-
rewriter.replaceOpWithNewOp<IREE::TensorExt::DispatchTensorLoadOp>(
158-
loadOp, adaptor.getSource(), loadOp.getSourceDims(), newOffsets,
159-
extractSizes, newStrides);
160-
return success();
161-
}
162-
};
163-
164-
/// Pattern to convert `iree_tensor_ext.dispatch.tensor.store` operation when
165-
/// materializing the encoding.
166-
struct MaterializeFlowDispatchTensorStoreOp final
167-
: OpConversionPattern<IREE::TensorExt::DispatchTensorStoreOp> {
168-
using Base::Base;
169-
170-
LogicalResult
171-
matchAndRewrite(IREE::TensorExt::DispatchTensorStoreOp storeOp,
172-
OpAdaptor adaptor,
173-
ConversionPatternRewriter &rewriter) const override {
174-
// Only handle operations where the store covers the entire
175-
// `!iree_tensor_ext.dispatch.tensor` type.
176-
if (!storeOp.isStoreToWholeTarget()) {
177-
return rewriter.notifyMatchFailure(storeOp, "unhandled partial stores");
178-
}
179-
180-
auto &typeConverter =
181-
*getTypeConverter<MaterializePadEncodingTypeConverter>();
182-
IREE::TensorExt::DispatchTensorType targetType = storeOp.getTargetType();
183-
auto boundTensorType = cast<RankedTensorType>(targetType.getBoundType());
184-
if (!typeConverter.hasNonZeroPadding(boundTensorType)) {
185-
// Let the Nop pattern handle this.
186-
return rewriter.notifyMatchFailure(storeOp, "no padding applied");
187-
}
188-
189-
IREE::TensorExt::DispatchTensorType newTargetType =
190-
typeConverter.convertType<IREE::TensorExt::DispatchTensorType>(
191-
targetType);
192-
RankedTensorType paddedType = newTargetType.asRankedTensorType();
193-
194-
Location loc = storeOp.getLoc();
195-
SmallVector<OpFoldResult> offsets(paddedType.getRank(),
196-
rewriter.getIndexAttr(0));
197-
SmallVector<OpFoldResult> strides(paddedType.getRank(),
198-
rewriter.getIndexAttr(1));
199-
SmallVector<OpFoldResult> sizes =
200-
tensor::getMixedSizes(rewriter, loc, adaptor.getValue());
201-
rewriter.replaceOpWithNewOp<IREE::TensorExt::DispatchTensorStoreOp>(
202-
storeOp, adaptor.getValue(), adaptor.getTarget(),
203-
adaptor.getTargetDims(), offsets, sizes, strides);
204-
return success();
205-
}
206-
};
207-
208119
/// Pattern to convert `iree_tensor_ext.dispatch.tensor.store` operation when
209120
/// materializing the encoding. We can not reuse the existing one because it
210121
/// does not transform new dynamic dimension through interface. The other
@@ -293,9 +204,7 @@ struct MaterializeEncodingIntoPaddingPass final
293204
// with the exception of a few ops that have to account for padding.
294205
// We add custom patterns with much higher priority to run before the
295206
// equivalent 'Nop' patterns.
296-
materializeEncodingPattern.add<MaterializeFlowDispatchTensorLoadOp,
297-
MaterializeFlowDispatchTensorStoreOp,
298-
MaterializeInterfaceBindingEncoding>(
207+
materializeEncodingPattern.add<MaterializeInterfaceBindingEncoding>(
299208
typeConverter, context, PatternBenefit{100});
300209

301210
if (failed(applyPartialConversion(operation, target,

‎compiler/src/iree/compiler/Codegen/ExternalInterfaces/GPUEncodingExternalModels.cpp‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -574,6 +574,24 @@ struct GPULayoutResolverAttr final
574574
struct GPUPadEncodingLayoutMaterializerAttr final
575575
: IREE::Encoding::LayoutMaterializerAttr::ExternalModel<
576576
GPUPadEncodingLayoutMaterializerAttr, GPUPaddingResolverAttr> {
577+
578+
LogicalResult getOffsetsSizesStrides(
579+
Attribute attr, OpBuilder &builder, Location loc,
580+
IREE::TensorExt::DispatchTensorType type, ValueRange dynamicDims,
581+
ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
582+
ArrayRef<OpFoldResult> strides, SmallVectorImpl<OpFoldResult> &newOffsets,
583+
SmallVectorImpl<OpFoldResult> &newSizes,
584+
SmallVectorImpl<OpFoldResult> &newStrides) const {
585+
auto boundType = dyn_cast<RankedTensorType>(type.getBoundType());
586+
if (!boundType || !boundType.getEncoding()) {
587+
return failure();
588+
}
589+
newSizes = getMixedValues(boundType.getShape(), dynamicDims, builder);
590+
newOffsets.resize(newSizes.size(), builder.getIndexAttr(0));
591+
newStrides.resize(newSizes.size(), builder.getIndexAttr(1));
592+
return success();
593+
}
594+
577595
Operation *lowerOp(Attribute attr, OpBuilder &b, Operation *op,
578596
TypeRange convertedResTypes,
579597
ValueRange convertedOperands) const {

‎compiler/src/iree/compiler/Codegen/ExternalInterfaces/Utils.h‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -128,17 +128,22 @@ struct EncodingLayoutMaterializerAttrExternalModelBase
128128
ArrayRef<OpFoldResult> strides, SmallVectorImpl<OpFoldResult> &newOffsets,
129129
SmallVectorImpl<OpFoldResult> &newSizes,
130130
SmallVectorImpl<OpFoldResult> &newStrides) const {
131-
auto layoutAttr = cast<IREE::Encoding::LayoutMaterializerAttr>(attr);
131+
auto boundType = dyn_cast<RankedTensorType>(type.getBoundType());
132+
if (!boundType || !boundType.getEncoding()) {
133+
return failure();
134+
}
135+
132136
// Only handle cases where the slice spans the whole
133137
// `!iree_tensor_ext.dispatch.tensor` type.
134138
// TODO(jornt): Enable partial slices.
135139
if (!type.doesSliceSpanWholeTensor(dynamicDims, offsets, sizes, strides)) {
136140
return failure();
137141
}
138-
auto boundTensorType = cast<RankedTensorType>(type.getBoundType());
142+
143+
auto layoutAttr = cast<IREE::Encoding::LayoutMaterializerAttr>(attr);
139144
IREE::Codegen::MaterializeEncodingInfo encodingInfo =
140-
getEncodingInfoFromLayout(boundTensorType, layoutAttr);
141-
newSizes = getMixedValues(boundTensorType.getShape(), dynamicDims, builder);
145+
getEncodingInfoFromLayout(boundType, layoutAttr);
146+
newSizes = getMixedValues(boundType.getShape(), dynamicDims, builder);
142147
FailureOr<SmallVector<OpFoldResult>> convertedMixedSizes =
143148
getPackedDimsForDispatchTensorImpl(builder, loc, type, dynamicDims,
144149
layoutAttr, encodingInfo);

0 commit comments

Comments
 (0)