[DT] Implement MaterializeInterfaceBindingEncoding with interface methods. (#22467)

The revision implements `LayoutMaterializerAttr::convertType` for
GPUPaddingResolverAttr; updates the
`MaterializeInterfaceBindingEncoding` pattern to use interface methods
to materialize the bindings.

It also fixes a bug in `@dynamic_set_zero_pad_encoding_and_store` test
that misses dynamic dimensions in binding ops.

It is a step towards #20160

---------

Signed-off-by: hanhanW <[email protected]>

Author: hanhanW

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

@@ -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));
   });
 }
 

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

@@ -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> {
@@ -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");

compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingPatterns.cpp

@@ -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.
@@ -175,44 +152,60 @@ struct MaterializeInterfaceBindingEncoding
   matchAndRewrite(IREE::HAL::InterfaceBindingSubspanOp subspanOp,
                   OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    auto resultType = llvm::dyn_cast<IREE::TensorExt::DispatchTensorType>(
+    auto origResultType = 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 =
+        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 = getTypeConverter<MaterializeEncodingTypeConverter>();
+    auto convertedResultType =
+        typeConverter->convertType<IREE::TensorExt::DispatchTensorType>(
+            origResultType);
+    if (!convertedResultType) {
+      return rewriter.notifyMatchFailure(subspanOp,
+                                         "expected converted result type to be "
+                                         "!iree_tensor_ext.dispatch.tensor");
+    }
+    if (origResultType == convertedResultType) {
+      return rewriter.notifyMatchFailure(
+          subspanOp, "DispatchTensorType type already valid");
     }
 
-    auto typeConverter = getTypeConverter<MaterializeEncodingTypeConverter>();
     // 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();
   }
 };

compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_into_padding.mlir

@@ -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>

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

@@ -580,10 +580,83 @@ 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) {
+    return nullptr;
+  }
+  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 type) {
+          // By the definition, the final converted type is the same tensor type
+          // without encodings.
+          return type.dropEncoding();
+        })
+        .Case([&](IREE::TensorExt::DispatchTensorType 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; });
+  }
+
   LogicalResult getOffsetsSizesStrides(
       Attribute attr, OpBuilder &builder, Location loc,
       IREE::TensorExt::DispatchTensorType type, ValueRange dynamicDims,