[DT] Implement LayoutMaterializerAttr for identity resolver. (#22337)

hanhanW · web-flow · commit 3815582bbd42 · 2025-10-21T10:12:23.000-07:00
The identity resolver is one of the basic encoding resolvers that drops encodings. The revision implements the interface and retires the IREE::Codegen::EncodingNopLayoutAttr. Fixes #20846 --------- Signed-off-by: hanhanW <hanhan0912@gmail.com>
diff --git a/compiler/src/iree/compiler/Codegen/Common/MaterializeEncoding.cpp b/compiler/src/iree/compiler/Codegen/Common/MaterializeEncoding.cpp
@@ -7,12 +7,10 @@
 #include "iree/compiler/Codegen/Common/EncodingUtils.h"
 #include "iree/compiler/Codegen/Common/Transforms.h"
 #include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.h"
-#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
 #include "iree/compiler/Codegen/Utils/GPUUtils.h"
-#include "iree/compiler/Codegen/Utils/Utils.h"
 #include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
 #include "iree/compiler/Dialect/HAL/Analysis/DeviceAnalysis.h"
 #include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
@@ -98,10 +96,10 @@ materializeFuncOpEncodings(FunctionOpInterface funcOp,
             IREE::GPU::GPUEncodingResolverAttr::get(
                 ctx, DictionaryAttr::get(ctx, configItems)));
       }
-      LDBG() << "Select EncodingNopLayoutAttr attribute as the layout "
+      LDBG() << "Select IdentityResolverAttr attribute as the layout "
                 "attribute (Encoding resolver unknown or unsupported).";
       return cast<IREE::Encoding::LayoutMaterializerAttr>(
-          IREE::Codegen::EncodingNopLayoutAttr::get(ctx));
+          IREE::Encoding::IdentityResolverAttr::get(ctx));
     };
 
     // The layoutAttr should come in without any target info attached to it,
diff --git a/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoNop.cpp b/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoNop.cpp
@@ -7,7 +7,6 @@
 #include "iree/compiler/Codegen/Common/EncodingUtils.h"
 #include "iree/compiler/Codegen/Common/PassUtils.h"
 #include "iree/compiler/Codegen/Common/Passes.h"
-#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
@@ -38,7 +37,7 @@ struct MaterializeEncodingIntoNopPass final
 
     RewritePatternSet materializeEncodingPattern(context);
     auto layoutAttr = cast<IREE::Encoding::LayoutMaterializerAttr>(
-        IREE::Codegen::EncodingNopLayoutAttr::get(context));
+        IREE::Encoding::IdentityResolverAttr::get(context));
     MaterializeEncodingTypeConverter typeConverter(layoutAttr);
     MaterializeEncodingConversionTarget target(*context);
     populateMaterializeEncodingPatterns(materializeEncodingPattern, target,
diff --git a/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoPadding.cpp b/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoPadding.cpp
@@ -8,7 +8,6 @@
 #include "iree/compiler/Codegen/Common/EncodingUtils.h"
 #include "iree/compiler/Codegen/Common/PassUtils.h"
 #include "iree/compiler/Codegen/Common/Passes.h"
-#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
 #include "iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.h"
@@ -191,7 +190,7 @@ struct MaterializeEncodingIntoPaddingPass final
           resolverAttr.cloneWithSimplifiedConfig(targetConfig));
     } else {
       layoutAttr = cast<IREE::Encoding::LayoutMaterializerAttr>(
-          IREE::Codegen::EncodingNopLayoutAttr::get(context));
+          IREE::Encoding::IdentityResolverAttr::get(context));
     }
 
     RewritePatternSet materializeEncodingPattern(context);
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_for_iree_ops.mlir b/compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_for_iree_ops.mlir
@@ -1,5 +1,78 @@
 // RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-materialize-device-encoding))" --split-input-file %s | FileCheck %s
 
+//----------------------------------------------------------------------------//
+// Test suite using generic encoding resolvers, that are defined in Encoding
+// dialect.
+//----------------------------------------------------------------------------//
+
+#pipeline_layout = #hal.pipeline.layout<constants = 3, bindings = [
+  #hal.pipeline.binding<storage_buffer>,
+  #hal.pipeline.binding<storage_buffer>,
+  #hal.pipeline.binding<storage_buffer>
+]>
+#map = affine_map<(d0, d1, d2) -> (d0, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
+#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
+#encoding_lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+#encoding_rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+#encoding_result = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], iteration_sizes = [?, ?, ?]>
+func.func @matmul_lowering_f32f32f32_identity_resolver() attributes {
+  hal.executable.target = #hal.executable.target<"llvm-cpu", "whatever", {iree.encoding.resolver = #iree_encoding.identity_resolver<>}>
+} {
+  %c0 = arith.constant 0 : index
+  %M = hal.interface.constant.load layout(#pipeline_layout) ordinal(0) : index
+  %N = hal.interface.constant.load layout(#pipeline_layout) ordinal(1) : index
+  %K = hal.interface.constant.load layout(#pipeline_layout) ordinal(2) : index
+  %0 = hal.interface.binding.subspan layout(#pipeline_layout) binding(0) alignment(64) offset(%c0)
+      : !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x?xf32, #encoding_lhs>>{%M, %K}
+  %1 = hal.interface.binding.subspan layout(#pipeline_layout) binding(1) alignment(64) offset(%c0)
+      : !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x?xf32, #encoding_rhs>>{%K, %N}
+  %2 = hal.interface.binding.subspan layout(#pipeline_layout) binding(2) alignment(64) offset(%c0)
+      : !iree_tensor_ext.dispatch.tensor<readwrite:tensor<?x?xf32, #encoding_result>>{%M, %N}
+  %3 = iree_tensor_ext.dispatch.tensor.load %0, offsets = [0, 0], sizes = [%M, %K], strides = [1, 1]
+      : !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x?xf32, #encoding_lhs>>{%M, %K}
+      -> tensor<?x?xf32, #encoding_lhs>
+  %4 = iree_tensor_ext.dispatch.tensor.load %1, offsets = [0, 0], sizes = [%K, %N], strides = [1, 1]
+      : !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x?xf32, #encoding_rhs>>{%K, %N}
+      -> tensor<?x?xf32, #encoding_rhs>
+  %5 = iree_tensor_ext.dispatch.tensor.load %2, offsets = [0, 0], sizes = [%M, %N], strides = [1, 1]
+      : !iree_tensor_ext.dispatch.tensor<readwrite:tensor<?x?xf32, #encoding_result>>{%M, %N}
+      -> tensor<?x?xf32, #encoding_result>
+  %6 = linalg.matmul
+      ins(%3, %4 : tensor<?x?xf32, #encoding_lhs>,
+                   tensor<?x?xf32, #encoding_rhs>)
+      outs(%5 : tensor<?x?xf32, #encoding_result>)
+      -> tensor<?x?xf32, #encoding_result>
+  iree_tensor_ext.dispatch.tensor.store %6, %2, offsets = [0, 0], sizes = [%M, %N], strides = [1, 1]
+      : tensor<?x?xf32, #encoding_result>
+      -> !iree_tensor_ext.dispatch.tensor<readwrite:tensor<?x?xf32, #encoding_result>>{%M, %N}
+  return
+}
+// CHECK-LABEL: func @matmul_lowering_f32f32f32_identity_resolver()
+//   CHECK-DAG:   %[[C0:.+]] = arith.constant 0 : index
+//   CHECK-DAG:   %[[M:.+]] = hal.interface.constant.load layout(#pipeline_layout) ordinal(0)
+//   CHECK-DAG:   %[[N:.+]] = hal.interface.constant.load layout(#pipeline_layout) ordinal(1)
+//   CHECK-DAG:   %[[K:.+]] = hal.interface.constant.load layout(#pipeline_layout) ordinal(2)
+//       CHECK:   %[[LHS_BINDING:.+]] = hal.interface.binding.subspan layout({{.+}}) binding(0)
+//  CHECK-SAME:       !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x?xf32>>{%[[M]], %[[K]]}
+//       CHECK:   %[[RHS_BINDING:.+]] = hal.interface.binding.subspan layout({{.+}}) binding(1)
+//  CHECK-SAME:       !iree_tensor_ext.dispatch.tensor<readonly:tensor<?x?xf32>>{%[[K]], %[[N]]}
+//       CHECK:   %[[OUTS_BINDING:.+]] = hal.interface.binding.subspan layout({{.+}}) binding(2)
+//  CHECK-SAME:       !iree_tensor_ext.dispatch.tensor<readwrite:tensor<?x?xf32>>{%[[M]], %[[N]]}
+//       CHECK:   %[[LHS:.+]] = iree_tensor_ext.dispatch.tensor.load %[[LHS_BINDING]]
+//  CHECK-SAME:       offsets = [0, 0], sizes = [%[[M]], %[[K]]], strides = [1, 1]
+//       CHECK:   %[[RHS:.+]] = iree_tensor_ext.dispatch.tensor.load %[[RHS_BINDING]]
+//  CHECK-SAME:       offsets = [0, 0], sizes = [%[[K]], %[[N]]], strides = [1, 1]
+//       CHECK:   %[[OUTS:.+]] = iree_tensor_ext.dispatch.tensor.load %[[OUTS_BINDING]]
+//  CHECK-SAME:       offsets = [0, 0], sizes = [%[[M]], %[[N]]], strides = [1, 1]
+//       CHECK:   %[[RES:.+]] = linalg.matmul
+//  CHECK-SAME:       ins(%[[LHS]], %[[RHS]] :
+//  CHECK-SAME:       outs(%[[OUTS]] :
+//       CHECK:   iree_tensor_ext.dispatch.tensor.store %[[RES]], %[[OUTS_BINDING]]
+//  CHECK-SAME:       offsets = [0, 0], sizes = [%[[M]], %[[N]]], strides = [1, 1]
+
+// -----
+
 //----------------------------------------------------------------------------//
 // Test suite using CPU encoding resolvers.
 //----------------------------------------------------------------------------//
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.td b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.td
@@ -424,21 +424,6 @@ def IREECodegen_ExportConfig : AttrDef<IREECodegen_Dialect, "ExportConfig", []>
   let genVerifyDecl = 1;
 }
 
-//===---------------------------------------------------------------------===//
-// iree_codegen.encoding_nop_layout
-//===---------------------------------------------------------------------===//
-
-def IREECodegen_EncodingNopLayoutAttr  :
-    AttrDef<IREECodegen_Dialect, "EncodingNopLayout"> {
-  let mnemonic = "encoding_nop_layout";
-  let summary = [{An attribute with implementation that treats encoding as nop.}];
-  let description = [{
-    An attribute that implements the interface methods that discards the
-    encodings. It can be a default attribute when a backend does not implement
-    encoding details.
-  }];
-}
-
 //===---------------------------------------------------------------------===//
 // iree_codegen.rotate_rows
 //===---------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/BUILD.bazel b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/BUILD.bazel
@@ -16,14 +16,12 @@ iree_compiler_cc_library(
     name = "ExternalModels",
     srcs = [
         "CPUEncodingExternalModels.cpp",
-        "CodegenExternalModels.cpp",
         "GPUEncodingExternalModels.cpp",
         "Interfaces.cpp",
         "Utils.cpp",
     ],
     hdrs = [
         "CPUEncodingExternalModels.h",
-        "CodegenExternalModels.h",
         "GPUEncodingExternalModels.h",
         "Interfaces.h",
         "Utils.h",
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CMakeLists.txt
@@ -15,13 +15,11 @@ iree_cc_library(
     ExternalModels
   HDRS
     "CPUEncodingExternalModels.h"
-    "CodegenExternalModels.h"
     "GPUEncodingExternalModels.h"
     "Interfaces.h"
     "Utils.h"
   SRCS
     "CPUEncodingExternalModels.cpp"
-    "CodegenExternalModels.cpp"
     "GPUEncodingExternalModels.cpp"
     "Interfaces.cpp"
     "Utils.cpp"
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CodegenExternalModels.cpp b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CodegenExternalModels.cpp
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CodegenExternalModels.h b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CodegenExternalModels.h
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Interfaces.cpp b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Interfaces.cpp
@@ -7,13 +7,11 @@
 #include "iree/compiler/Codegen/ExternalInterfaces/Interfaces.h"
 
 #include "iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.h"
-#include "iree/compiler/Codegen/ExternalInterfaces/CodegenExternalModels.h"
 #include "iree/compiler/Codegen/ExternalInterfaces/GPUEncodingExternalModels.h"
 
 namespace mlir::iree_compiler {
 
 void registerCodegenExternalInterfaces(DialectRegistry &registry) {
-  IREE::Codegen::registerCodegenExternalModels(registry);
   IREE::CPU::registerCPUEncodingExternalModels(registry);
   IREE::GPU::registerGPUEncodingExternalModels(registry);
 }
diff --git a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp
@@ -524,6 +524,53 @@ Attribute IdentityResolverAttr::getLayout(RankedTensorType type) const {
   return Encoding::IdentityAttr::get(getContext());
 }
 
+Type IdentityResolverAttr::convertType(Type type) const {
+  using IREE::TensorExt::DispatchTensorType;
+  return TypeSwitch<Type, Type>(type)
+      .Case<RankedTensorType>([&](auto rankedTensorType) {
+        return rankedTensorType.dropEncoding();
+      })
+      .Case<DispatchTensorType>([&](auto dispatchTensorType) {
+        auto boundType =
+            dyn_cast<RankedTensorType>(dispatchTensorType.getBoundType());
+        if (!boundType || !boundType.getEncoding()) {
+          return dispatchTensorType;
+        }
+        Type convertedBoundType = this->convertType(boundType);
+        return DispatchTensorType::get(dispatchTensorType.getAccess(),
+                                       convertedBoundType);
+      })
+      .Default([&](auto concreteType) { return concreteType; });
+}
+
+LogicalResult IdentityResolverAttr::getOffsetsSizesStrides(
+    OpBuilder &builder, Location loc, IREE::TensorExt::DispatchTensorType type,
+    ValueRange dynamicDims, ArrayRef<OpFoldResult> offsets,
+    ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides,
+    SmallVectorImpl<OpFoldResult> &newOffsets,
+    SmallVectorImpl<OpFoldResult> &newSizes,
+    SmallVectorImpl<OpFoldResult> &newStrides) const {
+  // Only handle cases where the slice spans the whole
+  // `!iree_tensor_ext.dispatch.tensor` type.
+  // TODO(hanchung): Enable partial slices. It was copied from pattern's
+  // implementaion, i.e., the users, and it can be dropped after we move the
+  // checks to the interface implementations.
+  if (!type.doesSliceSpanWholeTensor(dynamicDims, offsets, sizes, strides)) {
+    return failure();
+  }
+  auto boundTensorType = cast<RankedTensorType>(type.getBoundType());
+  newSizes = getMixedValues(boundTensorType.getShape(), dynamicDims, builder);
+  newOffsets.resize(newSizes.size(), builder.getIndexAttr(0));
+  newStrides.resize(newSizes.size(), builder.getIndexAttr(1));
+  return success();
+}
+
+Operation *IdentityResolverAttr::lowerOp(OpBuilder &b, Operation *op,
+                                         TypeRange convertedResTypes,
+                                         ValueRange convertedOperands) const {
+  return clone(b, op, convertedResTypes, convertedOperands);
+}
+
 //===---------------------------------------------------------------------===//
 // iree_encoding.unsupported_resolver
 //===---------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.td b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.td
@@ -269,6 +269,11 @@ def IdentityResolverAttr :
       DeclareAttrInterfaceMethods<IREEEncoding_LayoutResolverAttr, [
         "cloneWithSimplifiedConfig",
         "getLayout",
+      ]>,
+      DeclareAttrInterfaceMethods<IREEEncoding_LayoutMaterializerAttr, [
+        "convertType",
+        "getOffsetsSizesStrides",
+        "lowerOp",
       ]>
     ]> {
   let mnemonic = "identity_resolver";
