[Stream] Attach layouts to tensor ops in encoding specialization pass. (#19649)
The revision attaches the layouts to Stream tensor ops, e.g.,
stream.tensor.sizeof, etc. Each layout is wrapped with an attribute
where the attribute should implement the `calculateStorageSizeInBytes`
interface method. Because only the device/target knows the details. It
provides a way to materialize the serialized layouts into actual storage
size.
To teach the pass about layouts, there are few core changes happening in
the revision.
1. The revision introduces `cloneWithSimplifiedConfig` interface method
to EncodingLayoutAttrInterface. Each backend encoding attribute can
filter needed configurations when making layout resolvers in
HALAffinityAnalysisDialectInterface. This is needed because attributes
are not mutable. We always need to create a new attribute if we want to
update the parameters. The main benifit is that it makes the IR dump
much less verbose.
2. The revision also introduces `getLayout` interface method, which
returns the encoded layouts into encoding attributes. The CPU and VMVX
encoding attributes implement it in the revision, and it returns the
serialized MaterializeEncodingInfo struct as a dictionary attribute. See
below example for more details.
```
... layouts = [#iree_cpu.vmvx_encoding_layout<configuration =
{encoding_info =
{innerDimsPos = [0, 1],
innerTileSizes = [8, 4],
outerDimsPerm = [0, 1]}}>
]
```
---------
Signed-off-by: hanhanW <hanhan0912@gmail.com>diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/BUILD.bazel b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/BUILD.bazel
index 640e1ef..ec9b711 100644
--- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/BUILD.bazel
@@ -18,11 +18,13 @@
"CPUEncodingExternalModels.cpp",
"GPUEncodingExternalModels.cpp",
"Interfaces.cpp",
+ "Utils.cpp",
],
hdrs = [
"CPUEncodingExternalModels.h",
"GPUEncodingExternalModels.h",
"Interfaces.h",
+ "Utils.h",
],
deps = [
"//compiler/src/iree/compiler/Codegen/Dialect/CPU/IR:IREECPUDialect",
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CMakeLists.txt
index 46d0627..3f14e77 100644
--- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CMakeLists.txt
@@ -17,10 +17,12 @@
"CPUEncodingExternalModels.h"
"GPUEncodingExternalModels.h"
"Interfaces.h"
+ "Utils.h"
SRCS
"CPUEncodingExternalModels.cpp"
"GPUEncodingExternalModels.cpp"
"Interfaces.cpp"
+ "Utils.cpp"
DEPS
LLVMSupport
MLIRIR
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
index 3461a29..187cef7 100644
--- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
@@ -34,13 +34,14 @@
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenTypes.h"
#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h"
+#include "iree/compiler/Codegen/ExternalInterfaces/Utils.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
#include "llvm/Support/Debug.h"
#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
-#define DEBUG_TYPE "iree-gpu-encoding-external-models"
+#define DEBUG_TYPE "iree-cpu-encoding-external-models"
namespace mlir::iree_compiler::IREE::CPU {
@@ -53,6 +54,17 @@
// Utilities.
//===----------------------------------------------------------------------===//
+/// Appends the NamedAttribute into `config` if there is a `name` NamedAttribute
+/// in the `dictAttr`.
+static void storeNamedAttrIfPresent(SmallVectorImpl<NamedAttribute> &config,
+ DictionaryAttr dictAttr, StringRef name) {
+ auto attr = dictAttr.getNamed(name);
+ if (!attr) {
+ return;
+ }
+ config.push_back(attr.value());
+}
+
static void transposeInPlace(MaterializeEncodingInfo &info) {
// Vector cases: nothing to do.
if (info.innerTileSizes.size() < 2) {
@@ -640,6 +652,25 @@
}
};
+struct CPUHostEncodingLayoutAttrInterface
+ : public IREE::Encoding::EncodingLayoutAttrInterface::ExternalModel<
+ CPUHostEncodingLayoutAttrInterface, CPUEncodingLayoutAttr> {
+ Attribute cloneWithSimplifiedConfig(Attribute attr,
+ DictionaryAttr config) const {
+ MLIRContext *ctx = attr.getContext();
+ SmallVector<NamedAttribute> configItems;
+ storeNamedAttrIfPresent(configItems, config, "cpu_features");
+ storeNamedAttrIfPresent(configItems, config, "target_triple");
+ return CPUEncodingLayoutAttr::get(ctx,
+ DictionaryAttr::get(ctx, configItems));
+ }
+
+ Attribute getLayout(Attribute attr, RankedTensorType type) const {
+ MLIRContext *ctx = attr.getContext();
+ return CPUEncodingLayoutAttr::get(ctx, getLayoutImpl(attr, type));
+ }
+};
+
//===----------------------------------------------------------------------===//
// Interface methods implementaion for iree_cpu.vmvx_encoding_layout.
//===----------------------------------------------------------------------===//
@@ -728,15 +759,35 @@
}
};
+struct VMVXHostEncodingLayoutAttrInterface
+ : public IREE::Encoding::EncodingLayoutAttrInterface::ExternalModel<
+ VMVXHostEncodingLayoutAttrInterface, VMVXEncodingLayoutAttr> {
+ Attribute cloneWithSimplifiedConfig(Attribute attr,
+ DictionaryAttr config) const {
+ MLIRContext *ctx = attr.getContext();
+ SmallVector<NamedAttribute> configItems;
+ storeNamedAttrIfPresent(configItems, config, "ukernels");
+ return VMVXEncodingLayoutAttr::get(ctx,
+ DictionaryAttr::get(ctx, configItems));
+ }
+
+ Attribute getLayout(Attribute attr, RankedTensorType type) const {
+ MLIRContext *ctx = attr.getContext();
+ return VMVXEncodingLayoutAttr::get(ctx, getLayoutImpl(attr, type));
+ }
+};
+
} // namespace
void registerCPUEncodingExternalModels(DialectRegistry ®istry) {
registry.addExtension(
+[](MLIRContext *ctx, IREE::CPU::IREECPUDialect *dialect) {
IREE::CPU::CPUEncodingLayoutAttr::attachInterface<
- CPUDeviceEncodingLayoutAttrInterface>(*ctx);
+ CPUDeviceEncodingLayoutAttrInterface,
+ CPUHostEncodingLayoutAttrInterface>(*ctx);
IREE::CPU::VMVXEncodingLayoutAttr::attachInterface<
- VMVXDeviceEncodingLayoutAttrInterface>(*ctx);
+ VMVXDeviceEncodingLayoutAttrInterface,
+ VMVXHostEncodingLayoutAttrInterface>(*ctx);
});
}
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Utils.cpp b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Utils.cpp
new file mode 100644
index 0000000..12542e2
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Utils.cpp
@@ -0,0 +1,28 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Codegen/ExternalInterfaces/Utils.h"
+
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/MLIRContext.h"
+
+namespace mlir::iree_compiler::IREE {
+using Codegen::MaterializeEncodingInfo;
+
+DictionaryAttr getLayoutImpl(Attribute attr, RankedTensorType type) {
+ MLIRContext *ctx = attr.getContext();
+ auto deviceLayoutAttr = cast<IREE::Codegen::LayoutAttrInterface>(attr);
+ const MaterializeEncodingInfo info = deviceLayoutAttr.getEncodingInfo(type);
+ auto strAttr = StringAttr::get(ctx, "encoding_info");
+ Attribute encodingInfoAttr =
+ IREE::Codegen::serializeEncodingInfo(attr.getContext(), info);
+ return DictionaryAttr::get(ctx, {NamedAttribute(strAttr, encodingInfoAttr)});
+}
+
+} // namespace mlir::iree_compiler::IREE
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Utils.h b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Utils.h
new file mode 100644
index 0000000..f4649d6
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Utils.h
@@ -0,0 +1,22 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_COMPILER_CODEGEN_EXTERNALINTERFACES_UTILS_H_
+#define IREE_COMPILER_CODEGEN_EXTERNALINTERFACES_UTILS_H_
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+
+namespace mlir::iree_compiler::IREE {
+
+/// Returns a dictionary attribute that contains the materialized encoding info,
+/// i.e., serialized MaterializeEncodingInfo struct.
+/// Requirement: `attr` must implement IREE::Codegen::LayoutAttrInterface.
+DictionaryAttr getLayoutImpl(Attribute attr, RankedTensorType type);
+
+} // namespace mlir::iree_compiler::IREE
+
+#endif // IREE_COMPILER_CODEGEN_EXTERNALINTERFACES_UTILSS_H_
diff --git a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp
index c615972..8552645 100644
--- a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp
+++ b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp
@@ -260,6 +260,9 @@
FailureOr<linalg::ContractionDimensions>
getEncodingContractionDims(EncodingAttr encoding) {
auto indexingMapsAttr = encoding.getUserIndexingMaps();
+ if (!indexingMapsAttr) {
+ return failure();
+ }
SmallVector<AffineMap> indexingMaps = llvm::map_to_vector(
indexingMapsAttr.getValue(), [](Attribute m) -> AffineMap {
return cast<AffineMapAttr>(m).getAffineMap();
diff --git a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingInterfaces.td b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingInterfaces.td
index 4dfa4b5..6eb994f 100644
--- a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingInterfaces.td
+++ b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingInterfaces.td
@@ -43,6 +43,48 @@
assert(false && "unimplemented interface method");
return {};
}]
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Returns the attribute with simplified configuration/layouts. Attribute
+ is immutable in MLIR concept. Different attributes can implement
+ attribute interface methods differently, and they can carry target
+ configuration (e.g., cpu features) for further lowering. However, some
+ configuration/parameters can be dropped as long as they are no longer
+ needed in the progressively lowering. This method provides a mechanism
+ for such attribute to drop the outdated paramters and makes IR dump less
+ verbose.
+ }],
+ /*retTy=*/"::mlir::Attribute",
+ /*methodName=*/"cloneWithSimplifiedConfig",
+ /*args=*/(ins
+ "::mlir::DictionaryAttr":$config
+ ),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ assert(false && "unimplemented interface method");
+ return {};
+ }]
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Returns the serialized layout, which is either common format or wrapped
+ by an attribute that implements EncodingLayoutAttrInterface interface.
+ If it is in common format (e.g., a regular tensor type), we can easily
+ calculate the storage size. Otherwise, we will need a hook from
+ external, and the hook can come from an attribute that implements the
+ interface.
+ }],
+ /*retTy=*/"::mlir::Attribute",
+ /*methodName=*/"getLayout",
+ /*args=*/(ins
+ "::mlir::RankedTensorType":$type
+ ),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ assert(false && "unimplemented interface method");
+ return {};
+ }]
>
];
}
diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/BUILD.bazel b/compiler/src/iree/compiler/Dialect/HAL/IR/BUILD.bazel
index 576e77d..26311de 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/IR/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/HAL/IR/BUILD.bazel
@@ -100,6 +100,7 @@
hdrs = ["HALDialect.h"],
deps = [
":IR",
+ "//compiler/src/iree/compiler/Dialect/Encoding/IR",
"//compiler/src/iree/compiler/Dialect/HAL:hal_imports",
"//compiler/src/iree/compiler/Dialect/HAL/Analysis",
"//compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM",
diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/HAL/IR/CMakeLists.txt
index e0b68bd..ae64963 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/IR/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/HAL/IR/CMakeLists.txt
@@ -79,6 +79,7 @@
MLIRParser
MLIRSCFDialect
MLIRTransformUtils
+ iree::compiler::Dialect::Encoding::IR
iree::compiler::Dialect::HAL::Analysis
iree::compiler::Dialect::HAL::Conversion::HALToVM
iree::compiler::Dialect::HAL::hal_imports
diff --git a/compiler/src/iree/compiler/Dialect/HAL/IR/HALDialect.cpp b/compiler/src/iree/compiler/Dialect/HAL/IR/HALDialect.cpp
index e28d08f..a4335a7 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/IR/HALDialect.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/IR/HALDialect.cpp
@@ -6,6 +6,7 @@
#include "iree/compiler/Dialect/HAL/IR/HALDialect.h"
+#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
#include "iree/compiler/Dialect/HAL/Analysis/DeviceAnalysis.h"
#include "iree/compiler/Dialect/HAL/Conversion/HALToVM/Patterns.h"
#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
@@ -134,8 +135,18 @@
SetVector<IREE::HAL::ExecutableTargetAttr> resultSet;
deviceAnalysis.gatherRequiredExecutableTargets(affinityAttr, op,
resultSet);
- // TODO(hanchung): Populate the EncodingLayoutAttr when it is ready.
- layoutAttrs.insert(resultSet.begin(), resultSet.end());
+ for (auto targetAttr : resultSet) {
+ Attribute result = targetAttr;
+ if (auto attr = targetAttr.getConfiguration().getNamed("encoding")) {
+ if (auto encodingLayoutAttr =
+ dyn_cast<IREE::Encoding::EncodingLayoutAttrInterface>(
+ attr->getValue())) {
+ result = encodingLayoutAttr.cloneWithSimplifiedConfig(
+ targetAttr.getConfiguration());
+ }
+ }
+ layoutAttrs.insert(result);
+ }
return success();
};
};
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/SpecializeEncodings.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/SpecializeEncodings.cpp
index b177bed..40571f3 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/SpecializeEncodings.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/SpecializeEncodings.cpp
@@ -84,9 +84,9 @@
IRRewriter rewriter(funcOp.getContext());
for (auto affinityOp : candidates) {
auto affinityAttr = affinityOp.getAffinityAttr();
- SetVector<Attribute> layouts;
- if (failed(resolveLayoutAttr(affinityAttr, moduleOp, layouts))) {
- return affinityOp.emitError("failed on making layouts");
+ SetVector<Attribute> layoutResolvers;
+ if (failed(resolveLayoutAttr(affinityAttr, moduleOp, layoutResolvers))) {
+ return affinityOp.emitError("failed on making layout resolvers");
}
// Returns an updated encoding attribute if an encoding attribute is present
@@ -101,7 +101,17 @@
if (!encodingAttr) {
return std::nullopt;
}
- return encodingAttr.cloneWithLayouts(layouts.getArrayRef());
+ SmallVector<Attribute> layouts;
+ for (auto attr : layoutResolvers) {
+ auto encodingLayoutAttr =
+ dyn_cast<IREE::Encoding::EncodingLayoutAttrInterface>(attr);
+ if (!encodingLayoutAttr) {
+ layouts.push_back(attr);
+ continue;
+ }
+ layouts.push_back(encodingLayoutAttr.getLayout(rankedTensorType));
+ }
+ return encodingAttr.cloneWithLayouts(layouts);
};
// TODO(hanchung): Update other Stream operations.
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/specialize_encodings.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/specialize_encodings.mlir
index 1ae03e6..5fab86a 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/specialize_encodings.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/specialize_encodings.mlir
@@ -5,20 +5,31 @@
// encoding is updated that carries resolved layouts.
//------------------------------------------------------------------------------
-#executable_target_vmvx_bytecode_fb = #hal.executable.target<"vmvx", "vmvx-bytecode-fb", {encoding_layout = #iree_cpu.vmvx_encoding_layout<>}>
+#map0 = affine_map<(m, n, k) -> (m, k)>
+#map1 = affine_map<(m, n, k) -> (k, n)>
+#map2 = affine_map<(m, n, k) -> (m, n)>
+#executable_target_vmvx_bytecode_fb = #hal.executable.target<"vmvx", "vmvx-bytecode-fb", {encoding = #iree_cpu.vmvx_encoding_layout<>}>
+#executable_target_x86_64 = #hal.executable.target<"llvm-cpu", "xyz", {encoding = #iree_cpu.cpu_encoding_layout<>, target_triple="x86_64-xyz-xyz", cpu_features="+avx512f"}>
#device_target_local_0_ = #hal.device.target<"local", {ordinal = 0 : index}, [#executable_target_vmvx_bytecode_fb]> : !hal.device
-#encoding = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32]>
+#device_target_local_1_ = #hal.device.target<"local", {ordinal = 1 : index}, [#executable_target_x86_64]> : !hal.device
+#encoding = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map0, #map1, #map2]>
module {
util.global private @device_a = #device_target_local_0_
+ util.global private @device_b = #device_target_local_1_
- util.func public @tensor_sizeof(%d0: index, %d1: index) -> index {
- %size = stream.tensor.sizeof on(#hal.device.affinity<@device_a>) tensor<?x?xf32, #encoding>{%d0, %d1} : index
- util.return %size : index
+ util.func public @tensor_sizeof(%d0: index, %d1: index) -> (index, index) {
+ %size0 = stream.tensor.sizeof on(#hal.device.affinity<@device_a>) tensor<?x?xf32, #encoding>{%d0, %d1} : index
+ %size1 = stream.tensor.sizeof on(#hal.device.affinity<@device_b>) tensor<?x?xf32, #encoding>{%d0, %d1} : index
+ util.return %size0, %size1 : index, index
}
}
-// CHECK: #[[EXECUTABLE:.+]] = #hal.executable.target<"vmvx",
-// CHECK: #[[$ENCODING:.+]] = #iree_encoding.encoding
-// CHECK-SAME: layouts = [#[[EXECUTABLE]]]
+// CHECK: #[[$ENCODING0:.+]] = #iree_encoding.encoding
+// CHECK-SAME: #iree_cpu.vmvx_encoding_layout
+// CHECK-SAME: encoding_info = {innerDimsPos = [{{.+}}], innerTileSizes = [{{.+}}], outerDimsPerm = [{{.+}}]}
+// CHECK: #[[$ENCODING1:.+]] = #iree_encoding.encoding
+// CHECK-SAME: #iree_cpu.cpu_encoding_layout
+// CHECK-SAME: encoding_info = {innerDimsPos = [{{.+}}], innerTileSizes = [{{.+}}], outerDimsPerm = [{{.+}}]}
// CHECK-LABEL: util.func public @tensor_sizeof
-// CHECK: %[[RES:.+]] = stream.tensor.sizeof {{.+}} tensor<?x?xf32, #[[$ENCODING]]>
-// CHECK: return %[[RES]]
+// CHECK: %[[D0_RES:.+]] = stream.tensor.sizeof {{.+}} tensor<?x?xf32, #[[$ENCODING0]]>
+// CHECK: %[[D1_RES:.+]] = stream.tensor.sizeof {{.+}} tensor<?x?xf32, #[[$ENCODING1]]>
+// CHECK: return %[[D0_RES]], %[[D1_RES]]
