Skip to content

[DT] Implement MaterializeInterfaceBindingEncoding with interface methods. #22467

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 4 commits into from
Oct 30, 2025
+117 −190
Merged

[DT] Implement MaterializeInterfaceBindingEncoding with interface methods. #22467

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
bd6005d
[DT] Implement MaterializeInterfaceBindingEncoding with interface met…
hanhanW Oct 21, 2025
bf3fa4d
address comments
hanhanW Oct 30, 2025
ffc68d2
Merge branch 'main' into pad-resolver-convert-type
hanhanW Oct 30, 2025
cc10eea
drop & because we don't need it
hanhanW Oct 30, 2025
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
27 changes: 3 additions & 24 deletions compiler/src/iree/compiler/Codegen/Common/EncodingUtils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,32 +25,11 @@ MaterializeEncodingTypeConverter::MaterializeEncodingTypeConverter(
addConversion([](FloatType floatType) { return floatType; });
addConversion([](MemRefType memrefType) { return memrefType; });
addConversion([=](RankedTensorType type) {
// TODO(jornt): The isa<IREE::Encoding::PaddingAttr> check is
// needed because PaddingAttr is a serializable attribute, but it
// relies on its own type conversion for now. Once PaddingAttr
// implements `convertType`, this can be removed.
if (!isa<IREE::Encoding::PaddingAttr>(getLayoutAttr())) {
return cast<RankedTensorType>(getLayoutAttr().convertType(type));
}
return type.dropEncoding();
return cast<RankedTensorType>(getLayoutAttr().convertType(type));
});
addConversion([&](IREE::TensorExt::DispatchTensorType dispatchTensorType) {
auto boundType =
dyn_cast<RankedTensorType>(dispatchTensorType.getBoundType());
if (!boundType || !boundType.getEncoding()) {
return dispatchTensorType;
}
// TODO(jornt): The isa<IREE::Encoding::PaddingAttr> check is
// needed because PaddingAttr is a serializable attribute, but it
// relies on its own type conversion for now. Once PaddingAttr
// implements `convertType`, this can be removed.
if (!isa<IREE::Encoding::PaddingAttr>(getLayoutAttr())) {
return cast<IREE::TensorExt::DispatchTensorType>(
getLayoutAttr().convertType(dispatchTensorType));
}
Type convertedBoundType = convertType(boundType);
return IREE::TensorExt::DispatchTensorType::get(
dispatchTensorType.getAccess(), convertedBoundType);
return cast<IREE::TensorExt::DispatchTensorType>(
getLayoutAttr().convertType(dispatchTensorType));
});
}

Expand Down
120 changes: 1 addition & 119 deletions compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoPadding.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,116 +36,6 @@ namespace mlir::iree_compiler {
using namespace IREE::Encoding;

namespace {

// Returns the pad encoding layout, or nullptr if this is not the only layout or
// if there's no encoding at all.
static PaddingAttr getPadLayout(Attribute layoutAttr, RankedTensorType type) {
if (!type.getEncoding()) {
return nullptr;
}
auto encoding =
dyn_cast_or_null<IREE::Encoding::LayoutAttr>(type.getEncoding());
if (encoding) {
ArrayAttr layouts = encoding.getLayouts();
if (layouts.size() != 1) {
return nullptr;
}
return dyn_cast<PaddingAttr>(*layouts.begin());
}
Attribute resolvedEncoding =
cast<IREE::Encoding::LayoutResolverAttr>(layoutAttr).getLayout(type);
LLVM_DEBUG({
llvm::dbgs() << "Unresolved type: " << type << "\n";
llvm::dbgs() << "layoutAttr: " << layoutAttr << "\n";
llvm::dbgs() << "Resolved into: " << resolvedEncoding << "\n";
});
return dyn_cast<PaddingAttr>(resolvedEncoding);
}

// Returns a padded tensor type (without encoding) for tensor types with the pad
// encoding layout, or the same type for all other tensors.
static RankedTensorType getPaddedType(Attribute layoutAttr,
RankedTensorType type) {
PaddingAttr layout = getPadLayout(layoutAttr, type);
if (layout.isIdentityLayout()) {
return type.dropEncoding();
}

ArrayRef<int64_t> padding = layout.getPadding().asArrayRef();
auto newShape = llvm::to_vector_of<int64_t>(type.getShape());
for (auto [newDim, padValue] : llvm::zip_equal(newShape, padding)) {
assert((padValue == 0 || ShapedType::isStatic(newDim)) &&
"Padding dynamic dims not supported");
newDim += padValue;
}

return RankedTensorType::get(newShape, type.getElementType());
}

struct MaterializePadEncodingTypeConverter final
: MaterializeEncodingTypeConverter {
MaterializePadEncodingTypeConverter(
IREE::Encoding::LayoutMaterializerAttr layoutAttr)
: MaterializeEncodingTypeConverter(layoutAttr) {
addConversion([](RankedTensorType type) -> std::optional<RankedTensorType> {
// The type converter is designed for `padding` encoding
// attribute. By the definition, the final converted type is the same
// tensor type without encodings.
return type.dropEncoding();
});
addConversion([&](IREE::TensorExt::DispatchTensorType dispatchTensorType)
-> IREE::TensorExt::DispatchTensorType {
auto type = dyn_cast<RankedTensorType>(dispatchTensorType.getBoundType());
if (!type || !type.getEncoding()) {
return dispatchTensorType;
}
// The incoming bindings have the padded type, if `padding` is
// present.
if (getPadLayout(getLayoutAttr(), type)) {
type = getPaddedType(getLayoutAttr(), type);
}
return IREE::TensorExt::DispatchTensorType::get(
dispatchTensorType.getAccess(), type);
});
}

bool hasNonZeroPadding(RankedTensorType type) const {
PaddingAttr layout = getPadLayout(getLayoutAttr(), type);
return layout && !layout.isIdentityLayout();
}
};

/// Pattern to convert `iree_tensor_ext.dispatch.tensor.store` operation when
/// materializing the encoding. We can not reuse the existing one because it
/// does not transform new dynamic dimension through interface. The other
/// difference is that the converted type of the padding attribute is not as the
/// same as the tensor type that drops encoding.
/// TODO(#20160): Abstract new interface methods and collapse two patterns.
struct MaterializeInterfaceBindingEncoding final
: OpConversionPattern<IREE::HAL::InterfaceBindingSubspanOp> {
using Base::Base;

LogicalResult
matchAndRewrite(IREE::HAL::InterfaceBindingSubspanOp subspanOp,
OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto resultType = dyn_cast<IREE::TensorExt::DispatchTensorType>(
subspanOp.getResult().getType());
if (!resultType) {
return rewriter.notifyMatchFailure(
subspanOp,
"expected result type to be !iree_tensor_ext.dispatch.tensor");
}
auto newResultType = getTypeConverter()->convertType(resultType);
SmallVector<Value> newDynamicDims = subspanOp.getDynamicDims();
rewriter.replaceOpWithNewOp<IREE::HAL::InterfaceBindingSubspanOp>(
subspanOp, newResultType, subspanOp.getLayout(), subspanOp.getBinding(),
subspanOp.getByteOffset(), newDynamicDims, subspanOp.getAlignmentAttr(),
subspanOp.getDescriptorFlagsAttr());
return success();
}
};

struct MaterializeEncodingIntoPaddingPass final
: impl::MaterializeEncodingIntoPaddingPassBase<
MaterializeEncodingIntoPaddingPass> {
Expand Down Expand Up @@ -194,18 +84,10 @@ struct MaterializeEncodingIntoPaddingPass final
}

RewritePatternSet materializeEncodingPattern(context);
MaterializePadEncodingTypeConverter typeConverter(layoutAttr);
MaterializeEncodingTypeConverter typeConverter(layoutAttr);
MaterializeEncodingConversionTarget target(*context);
populateMaterializeEncodingPatterns(materializeEncodingPattern, target,
typeConverter);

// The majority of this conversion is based on the 'Nop' materialization,
// with the exception of a few ops that have to account for padding.
// We add custom patterns with much higher priority to run before the
// equivalent 'Nop' patterns.
materializeEncodingPattern.add<MaterializeInterfaceBindingEncoding>(
typeConverter, context, PatternBenefit{100});

if (failed(applyPartialConversion(operation, target,
std::move(materializeEncodingPattern)))) {
operation.emitOpError("materialization failed");
Expand Down
80 changes: 34 additions & 46 deletions compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingPatterns.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,29 +140,6 @@ lowerOpWithEncoding(RewriterBase &rewriter, tensor::EmptyOp emptyOp,
return newEmptyOp;
}

/// For `dispatchTensorType` that bind a `RankedTensorType` with encoding,
/// returns the dynamic dimensions of the materialized shape of the
/// `dispatchTensorType`. The dynamic dimensions of the `dispatchTensorType` are
/// provided in `dynamicDims`.
static FailureOr<SmallVector<Value>> getPackedDynamicDimsForDispatchTensor(
OpBuilder &builder, Location loc,
const MaterializeEncodingTypeConverter &typeConverter,
IREE::TensorExt::DispatchTensorType dispatchTensorType,
ValueRange dynamicDims) {
FailureOr<SmallVector<OpFoldResult>> convertedTargetShape =
typeConverter.getPackedDimsForDispatchTensor(
builder, loc, dispatchTensorType, dynamicDims);
if (failed(convertedTargetShape)) {
return failure();
}
SmallVector<int64_t> convertedStaticTargetShape;
SmallVector<Value> convertedDynamicTargetShape;
dispatchIndexOpFoldResults(convertedTargetShape.value(),
convertedDynamicTargetShape,
convertedStaticTargetShape);
return convertedDynamicTargetShape;
}

namespace {
/// Pattern to materialize the encoding for `hal.interface.binding.subspan`
/// operations.
Expand All @@ -175,45 +152,56 @@ struct MaterializeInterfaceBindingEncoding
matchAndRewrite(IREE::HAL::InterfaceBindingSubspanOp subspanOp,
OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto resultType = llvm::dyn_cast<IREE::TensorExt::DispatchTensorType>(
auto origResultType = llvm::dyn_cast<IREE::TensorExt::DispatchTensorType>(
subspanOp.getResult().getType());
if (!resultType) {
if (!origResultType) {
return rewriter.notifyMatchFailure(
subspanOp,
"expected result type to be !iree_tensor_ext.dispatch.tensor");
}
auto boundTensorType =
llvm::dyn_cast<RankedTensorType>(resultType.getBoundType());
if (!boundTensorType) {
auto origBoundTensorType =
llvm::dyn_cast<RankedTensorType>(origResultType.getBoundType());
if (!origBoundTensorType) {
return rewriter.notifyMatchFailure(
subspanOp, "bound type is not a RankedTensorType");
}

auto convertedBoundType = getTypeConverter()->convertType(boundTensorType);
if (convertedBoundType == boundTensorType) {
return rewriter.notifyMatchFailure(subspanOp, "bound type already valid");
}

auto *typeConverter = static_cast<const MaterializeEncodingTypeConverter *>(
getTypeConverter());
auto convertedResultType =
typeConverter->convertType<IREE::TensorExt::DispatchTensorType>(
origResultType);
if (origResultType == convertedResultType) {
return rewriter.notifyMatchFailure(
subspanOp, "DispatchTensorType type already valid");
}

// Get the dynamic dims of the target.
// TODO(hanchung): We only have getOffsetsSizesStrides interface method that
// handles all three together. It would be cleaner to have a separate method
// to get dynamic sizes only.
Location loc = subspanOp.getLoc();
SmallVector<Value> newDynamicDims = subspanOp.getDynamicDims();
FailureOr<SmallVector<Value>> convertedDynamicDims =
getPackedDynamicDimsForDispatchTensor(rewriter, loc, *typeConverter,
resultType,
subspanOp.getDynamicDims());
// Drop the encoding if the target does not support it.
if (succeeded(convertedDynamicDims)) {
newDynamicDims = convertedDynamicDims.value();
ValueRange origDynamicDims = subspanOp.getDynamicDims();
SmallVector<OpFoldResult> origSizes = getMixedValues(
origBoundTensorType.getShape(), origDynamicDims, rewriter);
SmallVector<OpFoldResult> origOffsets(origDynamicDims.size(),
rewriter.getIndexAttr(0));
SmallVector<OpFoldResult> origStrides(origDynamicDims.size(),
rewriter.getIndexAttr(1));
SmallVector<OpFoldResult> newOffsets, newSizes, newStrides;
if (failed(typeConverter->getOffsetsSizesStrides(
rewriter, loc, origResultType, origDynamicDims, origOffsets,
origSizes, origStrides, newOffsets, newSizes, newStrides))) {
return failure();
}

auto newResultType = IREE::TensorExt::DispatchTensorType::get(
resultType.getAccess(), convertedBoundType);
SmallVector<int64_t> newStaticDims;
SmallVector<Value> newDynamicDims;
dispatchIndexOpFoldResults(newSizes, newDynamicDims, newStaticDims);
rewriter.replaceOpWithNewOp<IREE::HAL::InterfaceBindingSubspanOp>(
subspanOp, newResultType, subspanOp.getLayout(), subspanOp.getBinding(),
subspanOp.getByteOffset(), newDynamicDims, subspanOp.getAlignmentAttr(),
subspanOp.getDescriptorFlagsAttr());
subspanOp, convertedResultType, subspanOp.getLayout(),
subspanOp.getBinding(), subspanOp.getByteOffset(), newDynamicDims,
subspanOp.getAlignmentAttr(), subspanOp.getDescriptorFlagsAttr());
return success();
}
};
Expand Down
4 changes: 2 additions & 2 deletions compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_into_padding.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,9 +113,9 @@ func.func @dynamic_set_zero_pad_encoding_and_store() {
%2 = hal.interface.constant.load layout(<constants = 2, bindings = [#binding_ro, #binding], flags = Indirect>) ordinal(1) : i32
%dynamic_sz = arith.index_castui %2 : i32 to index
%3 = hal.interface.binding.subspan layout(<constants = 2, bindings = [#binding_ro, #binding], flags = Indirect>) binding(0) alignment(64) offset(%1) flags("ReadOnly|Indirect")
: !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x2048xf16>>
: !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x2048xf16>>{%dynamic_sz}
%4 = hal.interface.binding.subspan layout(<constants = 2, bindings = [#binding_ro, #binding], flags = Indirect>) binding(1) alignment(64) offset(%c0) flags(Indirect)
: !iree_tensor_ext.dispatch.tensor<writeonly:tensor<?x2048xf16, #pad_encoding>>
: !iree_tensor_ext.dispatch.tensor<writeonly:tensor<?x2048xf16, #pad_encoding>>{%dynamic_sz}
%5 = iree_tensor_ext.dispatch.tensor.load %3, offsets = [0, 0], sizes = [%dynamic_sz, 2048], strides = [1, 1]
: !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x2048xf16>>{%dynamic_sz} -> tensor<?x2048xf16>
%6 = iree_encoding.set_encoding %5 : tensor<?x2048xf16> -> tensor<?x2048xf16, #encoding_mmt>
Expand Down
71 changes: 71 additions & 0 deletions compiler/src/iree/compiler/Codegen/ExternalInterfaces/GPUEncodingExternalModels.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -580,10 +580,81 @@ struct GPULayoutResolverAttr final
}
};

// Returns the pad encoding layout, or nullptr if this is not the only layout or
// if there's no encoding at all.
static IREE::Encoding::PaddingAttr
getPadLayout(IREE::Encoding::LayoutResolverAttr layoutAttr,
RankedTensorType type) {
if (!type.getEncoding()) {
return nullptr;
}
auto encoding =
dyn_cast_or_null<IREE::Encoding::LayoutAttr>(type.getEncoding());
if (encoding) {
ArrayAttr layouts = encoding.getLayouts();
if (layouts.size() != 1) {
return nullptr;
}
return dyn_cast<IREE::Encoding::PaddingAttr>(*layouts.begin());
}
Attribute resolvedEncoding = layoutAttr.getLayout(type);
LDBG() << "Unresolved type: " << type;
LDBG() << "layoutAttr: " << layoutAttr;
LDBG() << "Resolved into: " << resolvedEncoding;
return dyn_cast<IREE::Encoding::PaddingAttr>(resolvedEncoding);
}

// Returns a padded tensor type (without encoding) for tensor types with the pad
// encoding layout, or the same type for all other tensors.
static RankedTensorType
getPaddedType(IREE::Encoding::LayoutResolverAttr layoutAttr,
RankedTensorType type) {
IREE::Encoding::PaddingAttr layout = getPadLayout(layoutAttr, type);
if (layout.isIdentityLayout()) {
return type.dropEncoding();
}

ArrayRef<int64_t> padding = layout.getPadding().asArrayRef();
auto newShape = llvm::to_vector_of<int64_t>(type.getShape());
for (auto [newDim, padValue] : llvm::zip_equal(newShape, padding)) {
assert((padValue == 0 || ShapedType::isStatic(newDim)) &&
"Padding dynamic dims not supported");
newDim += padValue;
}

return RankedTensorType::get(newShape, type.getElementType());
}

struct GPUPadEncodingLayoutMaterializerAttr final
: IREE::Encoding::LayoutMaterializerAttr::ExternalModel<
GPUPadEncodingLayoutMaterializerAttr, GPUPaddingResolverAttr> {

Type convertType(Attribute attr, Type type) const {
auto layoutAttr = cast<IREE::Encoding::LayoutResolverAttr>(attr);
return TypeSwitch<Type, Type>(type)
.Case<RankedTensorType>([&](auto type) {
// By the definition, the final converted type is the same tensor type
// without encodings.
return type.dropEncoding();
})
.Case<IREE::TensorExt::DispatchTensorType>(
[&](auto dispatchTensorType) {
auto type =
dyn_cast<RankedTensorType>(dispatchTensorType.getBoundType());
if (!type || !type.getEncoding()) {
return dispatchTensorType;
}
// The incoming bindings have the padded type, if `padding` is
// present.
if (getPadLayout(layoutAttr, type)) {
type = getPaddedType(layoutAttr, type);
}
return IREE::TensorExt::DispatchTensorType::get(
dispatchTensorType.getAccess(), type);
})
.Default([&](Type type) { return type; });
Copy link
Member

Choose a reason for hiding this comment

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

Suggested change
.Default([&](Type type) { return type; });
.Default([](Type type) { return type; });

Copy link
Contributor Author

Choose a reason for hiding this comment

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

oh thanks, somehow I missed it.

}

LogicalResult getOffsetsSizesStrides(
Attribute attr, OpBuilder &builder, Location loc,
IREE::TensorExt::DispatchTensorType type, ValueRange dynamicDims,
Expand Down
Loading