[DT][VMVX] Implement VMVXEncodingLayoutAttr. (#19403)
The revision implements the VMVX encoding layout attribute for device
code with the new IREE::CPU dialect.
Some additional dialect registration is needed because it could create
the VMVXEncodingLayoutAttr attribute in the pass pipeline; the dialect
needs to be loaded. The main refactoring changes are:
- Move chooseMatmulTile to IREECodegen/Utils. We can internalize the
method once other CPU materialization logic is moved to their own
attribute implementation.
- Move the enumerateMatmulTilesVMVX method to
CPUEncodingExternalModels.cpp. It is only used by VMVX implementation.
On the Codegen utilitiy side (i.e., Codegen/Utils.[h|cpp]), the revision
adapts config query methods to use `Attribute`, and share the logic with
DictionaryAttr type. Previously, the configuration is wrapped into
IREE::HAL::ExecutableTarget attribute. Now we have DictionaryAttr
variants and we do not want to duplicate the implementation. Thus, the
functions take Attribute input and handle the cases in the
implementation.
On the data-tiling encoding materialization pass side, it creates the
VMVXEncodingLayoutAttr attribute (with the original target
configuration).
Note that the target configuration is an optional parameter and it is
expected to be used within pass scope, but not the final IR output.
No additional tests because they are covered by
vmvx_materialize_encoding.mlir.
---------
Signed-off-by: hanhanW <hanhan0912@gmail.com>
diff --git a/compiler/plugins/target/LLVMCPU/BUILD.bazel b/compiler/plugins/target/LLVMCPU/BUILD.bazel
index c086f54..fc25010 100644
--- a/compiler/plugins/target/LLVMCPU/BUILD.bazel
+++ b/compiler/plugins/target/LLVMCPU/BUILD.bazel
@@ -33,6 +33,7 @@
":StaticLibraryGenerator",
"//compiler/plugins/target/LLVMCPU/Builtins",
"//compiler/src/iree/compiler/Codegen/Common",
+ "//compiler/src/iree/compiler/Codegen/Dialect/CPU/IR:IREECPUDialect",
"//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
"//compiler/src/iree/compiler/Codegen/LLVMCPU",
"//compiler/src/iree/compiler/Codegen/Utils",
diff --git a/compiler/plugins/target/LLVMCPU/CMakeLists.txt b/compiler/plugins/target/LLVMCPU/CMakeLists.txt
index fea4eb8..9011c6d 100644
--- a/compiler/plugins/target/LLVMCPU/CMakeLists.txt
+++ b/compiler/plugins/target/LLVMCPU/CMakeLists.txt
@@ -54,6 +54,7 @@
MLIRTargetLLVMIRExport
MLIRTransformDialect
iree::compiler::Codegen::Common
+ iree::compiler::Codegen::Dialect::CPU::IR::IREECPUDialect
iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
iree::compiler::Codegen::LLVMCPU
iree::compiler::Codegen::Utils
diff --git a/compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp b/compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp
index ee8e256..c15b65c 100644
--- a/compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp
+++ b/compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp
@@ -15,6 +15,7 @@
#include "compiler/plugins/target/LLVMCPU/LibraryBuilder.h"
#include "compiler/plugins/target/LLVMCPU/LinkerTool.h"
#include "compiler/plugins/target/LLVMCPU/StaticLibraryGenerator.h"
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
#include "iree/compiler/Codegen/LLVMCPU/Passes.h"
#include "iree/compiler/Codegen/LLVMCPU/Utils.h"
@@ -218,6 +219,7 @@
// TODO: make inclusion of ArmNeon conditional?
// clang-format off
registry.insert<IREE::Codegen::IREECodegenDialect,
+ IREE::CPU::IREECPUDialect,
IREE::LinalgExt::IREELinalgExtDialect,
mlir::transform::TransformDialect,
pdl::PDLDialect,
diff --git a/compiler/plugins/target/VMVX/BUILD.bazel b/compiler/plugins/target/VMVX/BUILD.bazel
index 31d8b87..7b0cb11 100644
--- a/compiler/plugins/target/VMVX/BUILD.bazel
+++ b/compiler/plugins/target/VMVX/BUILD.bazel
@@ -23,6 +23,7 @@
"VMVXTarget.cpp",
],
deps = [
+ "//compiler/src/iree/compiler/Codegen/Dialect/CPU/IR:IREECPUDialect",
"//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
"//compiler/src/iree/compiler/Codegen/VMVX",
"//compiler/src/iree/compiler/Dialect/HAL/Target",
diff --git a/compiler/plugins/target/VMVX/CMakeLists.txt b/compiler/plugins/target/VMVX/CMakeLists.txt
index 7af0a88..f13273d 100644
--- a/compiler/plugins/target/VMVX/CMakeLists.txt
+++ b/compiler/plugins/target/VMVX/CMakeLists.txt
@@ -27,6 +27,7 @@
MLIRIR
MLIRPass
MLIRSupport
+ iree::compiler::Codegen::Dialect::CPU::IR::IREECPUDialect
iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
iree::compiler::Codegen::VMVX
iree::compiler::Dialect::HAL::Target
diff --git a/compiler/plugins/target/VMVX/VMVXTarget.cpp b/compiler/plugins/target/VMVX/VMVXTarget.cpp
index daba862..20a076e 100644
--- a/compiler/plugins/target/VMVX/VMVXTarget.cpp
+++ b/compiler/plugins/target/VMVX/VMVXTarget.cpp
@@ -4,6 +4,7 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
#include "iree/compiler/Codegen/VMVX/Passes.h"
#include "iree/compiler/Dialect/HAL/Target/Devices/LocalDevice.h"
@@ -77,9 +78,10 @@
}
void getDependentDialects(DialectRegistry ®istry) const override {
- registry.insert<IREE::Codegen::IREECodegenDialect, IREE::VM::VMDialect,
- IREE::VMVX::VMVXDialect,
- IREE::LinalgExt::IREELinalgExtDialect>();
+ registry
+ .insert<IREE::Codegen::IREECodegenDialect, IREE::CPU::IREECPUDialect,
+ IREE::VM::VMDialect, IREE::VMVX::VMVXDialect,
+ IREE::LinalgExt::IREELinalgExtDialect>();
}
IREE::VM::TargetOptions
@@ -232,8 +234,8 @@
}
void getDependentDialects(DialectRegistry ®istry) const override {
- registry
- .insert<IREE::Codegen::IREECodegenDialect, IREE::VMVX::VMVXDialect>();
+ registry.insert<IREE::Codegen::IREECodegenDialect,
+ IREE::CPU::IREECPUDialect, IREE::VMVX::VMVXDialect>();
}
void
diff --git a/compiler/src/iree/compiler/Codegen/Common/CPU/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/CPU/BUILD.bazel
index d6b3c96..f1053da 100644
--- a/compiler/src/iree/compiler/Codegen/Common/CPU/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Common/CPU/BUILD.bazel
@@ -56,6 +56,7 @@
":PassHeaders",
":PassesIncGen",
"//compiler/src/iree/compiler/Codegen/Common",
+ "//compiler/src/iree/compiler/Codegen/Dialect/CPU/IR:IREECPUDialect",
"//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
"//compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils",
"//compiler/src/iree/compiler/Codegen/Interfaces:UKernelOpInterface",
diff --git a/compiler/src/iree/compiler/Codegen/Common/CPU/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/CPU/CMakeLists.txt
index 81cefe7..75db95e 100644
--- a/compiler/src/iree/compiler/Codegen/Common/CPU/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Common/CPU/CMakeLists.txt
@@ -78,6 +78,7 @@
MLIRVectorTransforms
iree::builtins::ukernel::exported_bits
iree::compiler::Codegen::Common
+ iree::compiler::Codegen::Dialect::CPU::IR::IREECPUDialect
iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
iree::compiler::Codegen::Dialect::Codegen::Utils
iree::compiler::Codegen::Interfaces::UKernelOpInterface
diff --git a/compiler/src/iree/compiler/Codegen/Common/CPU/CPUMaterializeEncodings.cpp b/compiler/src/iree/compiler/Codegen/Common/CPU/CPUMaterializeEncodings.cpp
index a7c1f00..4d80343 100644
--- a/compiler/src/iree/compiler/Codegen/Common/CPU/CPUMaterializeEncodings.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/CPU/CPUMaterializeEncodings.cpp
@@ -6,6 +6,8 @@
#include "iree/compiler/Codegen/Common/CPU/Passes.h"
#include "iree/compiler/Codegen/Common/EncodingUtils.h"
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.h"
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h"
@@ -38,39 +40,9 @@
#define GEN_PASS_DEF_CPUMATERIALIZEHOSTENCODINGPASS
#include "iree/compiler/Codegen/Common/CPU/Passes.h.inc"
-// Enumerate tile sizes to choose from when no specific architecture is
-// targeted. For narrow-{M,N} cases, this only enumerates on narrow M. The
-// narrow-N cases are handled by transposition in chooseMatmulTile.
-static SmallVector<TileMxNxK>
-enumerateMatmulTilesVMVX(linalg::ContractionDimensions cDims,
- IREE::Encoding::EncodingAttr encoding,
- IREE::HAL::ExecutableTargetAttr target) {
- bool hasUkernelSupport = hasUkernel(target);
-
- // TODO(hanchung): The ukernel path does not support 3d
- // codegen.query_tile_sizes op, so we disable dynamic tile shapes for
- // batch_matmul. Also, they are not set up for narrow M/N matmul, so it is
- // disabled when it is the case.
- if (!cDims.batch.empty() || getMatmulNarrowDim(encoding)) {
- hasUkernelSupport = false;
- }
- if (hasUkernelSupport) {
- // VMVX+ukernel uses dynamic tile shapes.
- return {TileMxNxK{ShapedType::kDynamic, ShapedType::kDynamic,
- ShapedType::kDynamic}};
- }
-
- return {
- TileMxNxK{8, 8, 4}, // Some vaguely reasonable tile shape.
- TileMxNxK{4, 8, 4}, // Truncation of the above.
- TileMxNxK{2, 8, 4}, // Truncation of the above.
- TileMxNxK{1, 8, 4}, // Truncation of the above.
- };
-}
-
// Enumerate tile sizes to choose from on riscv32.
// For narrow-{M,N} cases, this only enumerates on narrow M. The narrow-N cases
-// are handled by transposition in chooseMatmulTile.
+// are handled by transposition in IREE::Codegen::chooseMatmulTile.
static SmallVector<TileMxNxK>
enumerateMatmulTileRiscv32(IREE::HAL::ExecutableTargetAttr target) {
if (hasUkernel(target)) {
@@ -87,7 +59,7 @@
// Enumerate tile sizes to choose from on arm64.
// For narrow-{M,N} cases, this only enumerates on narrow M. The narrow-N cases
-// are handled by transposition in chooseMatmulTile.
+// are handled by transposition in IREE::Codegen::chooseMatmulTile.
static SmallVector<TileMxNxK>
enumerateMatmulTileArm64(TypeRange elementTypes,
IREE::HAL::ExecutableTargetAttr target) {
@@ -178,7 +150,7 @@
// Enumerate tile sizes to choose from on x86-64.
// For narrow-{M,N} cases, this only enumerates on narrow M. The narrow-N cases
-// are handled by transposition in chooseMatmulTile.
+// are handled by transposition in IREE::Codegen::chooseMatmulTile.
static SmallVector<TileMxNxK>
enumerateMatmulTileX86_64(TypeRange elementTypes,
IREE::HAL::ExecutableTargetAttr target) {
@@ -291,114 +263,6 @@
return {};
}
-/// Returns the best TileMxNxK from `enumeratedTiles` pool. If the
-/// `hostDefinedUpperBound` is not empty, the chosen tile sizes can not be
-/// greater than the values.
-/// TODO(#16933): Remove `hostDefinedUpperBound` once we can propagate such
-/// information to host. For now, they are defined by host.
-static TileMxNxK
-chooseMatmulTile(ArrayRef<TileMxNxK> enumeratedTiles,
- IREE::Encoding::MatmulNarrowDim narrowDim,
- ArrayRef<int64_t> hostDefinedUpperBound = {}) {
- assert((hostDefinedUpperBound.empty() || hostDefinedUpperBound.size() >= 3) &&
- "expected hostDefinedUpperBound is empty or has upper bound for {M, "
- "N, K}");
- // Handle narrow-N by transposing to reduce to narrow-M. Note: the
- // enumeratedTiles currently only enumerate narrow-M cases.
- if (narrowDim.isN()) {
- SmallVector<int64_t> newHostDefinedUpperBound(hostDefinedUpperBound);
- std::swap(newHostDefinedUpperBound[0], newHostDefinedUpperBound[1]);
- narrowDim.dim = IREE::Encoding::MatmulNarrowDim::Dim::M;
- TileMxNxK tile =
- chooseMatmulTile(enumeratedTiles, narrowDim, newHostDefinedUpperBound);
- std::swap(tile.M, tile.N);
- return tile;
- }
- // Handle kDynamic: currently this is only used with VMVX, where there is only
- // one enumerated tile and it has all three M/N/K dimensions dynamic, so for
- // now we only support that. Generalize that as needed when more dynamic tile
- // sizes are used outside of VMVX, e.g. perhaps some day with Arm SVE. Decide
- // how to incorporate the handling of kDynamic in the cost-model evaluation
- // below to decide when to prefer a dynamic vs a static tile shape.
- for (auto tile : enumeratedTiles) {
- if (ShapedType::isDynamic(tile.M) || ShapedType::isDynamic(tile.N) ||
- ShapedType::isDynamic(tile.K)) {
- assert(enumeratedTiles.size() == 1);
- assert(ShapedType::isDynamic(tile.M) && ShapedType::isDynamic(tile.N) &&
- ShapedType::isDynamic(tile.K));
- return tile;
- }
- }
- // We're going to "rate" the enumerated tiles.
- struct RatedTileMxNxK : TileMxNxK {
- RatedTileMxNxK() {}
- RatedTileMxNxK(TileMxNxK tile) : TileMxNxK(tile) {}
- // Penalize tiles that are wider in the M dimension than matmulNarrowM.
- int64_t paddingPenalty = 0;
- // Favor larger tiles, as long as they still minimize paddingPenalty.
- int64_t productMxNxK = 0;
- };
- SmallVector<RatedTileMxNxK> ratedTiles;
- ratedTiles.reserve(enumeratedTiles.size());
- int64_t bestPaddingPenalty = INT64_MAX;
- int64_t mUB = INT64_MAX;
- int64_t nUB = INT64_MAX;
- int64_t kUB = INT64_MAX;
- if (!hostDefinedUpperBound.empty()) {
- mUB = hostDefinedUpperBound[0];
- nUB = hostDefinedUpperBound[1];
- kUB = hostDefinedUpperBound[2];
- }
- for (auto tile : enumeratedTiles) {
- if (tile.M > mUB || tile.N > nUB || tile.K > kUB) {
- LLVM_DEBUG(llvm::dbgs() << "[" << DEBUG_TYPE << "]: tile (";
- llvm::interleaveComma(
- ArrayRef<int64_t>{tile.M, tile.N, tile.K}, llvm::dbgs());
- llvm::dbgs()
- << ") is skipped because it is not valid for upper_bound (";
- llvm::interleaveComma(ArrayRef<int64_t>{mUB, nUB, kUB},
- llvm::dbgs());
- llvm::dbgs() << ")\n");
- continue;
- }
- RatedTileMxNxK ratedTile(tile);
- ratedTile.paddingPenalty = 0;
- // If we are choosing a tile for a narrow-M case, we want to minimize
- // padding along the M dimension.
- // The PowerOf2Ceil is so that we are OK with padding up to the next
- // power of two, we just try to avoid padding beyond that. For example,
- // if matmulNarrowM==7 and we have enumerated tiles with M=8,4,2,1, we
- // are OK with the tile that has M==8 even though it requires some padding.
- // Otherwise, we would be penalizing the tiles with M==8,4,2 and we would
- // end up selecting the vecmat tile (M==1) for that case!
- if (narrowDim) {
- ratedTile.paddingPenalty =
- std::max<int64_t>(tile.M - llvm::PowerOf2Ceil(narrowDim.size), 0);
- }
- ratedTile.productMxNxK = tile.M * tile.N * tile.K;
- ratedTiles.push_back(ratedTile);
-
- LLVM_DEBUG(llvm::dbgs() << "candidate: "; llvm::interleaveComma(
- ArrayRef<int64_t>{tile.M, tile.N, tile.K}, llvm::dbgs());
- llvm::dbgs() << " penalty:" << ratedTile.paddingPenalty << "\n");
-
- bestPaddingPenalty = std::min(bestPaddingPenalty, ratedTile.paddingPenalty);
- }
- RatedTileMxNxK bestRatedTile;
- for (auto ratedTile : ratedTiles) {
- // Choose only among tiles that minimize paddingPenalty. Among those,
- // maximize productMxNxK.
- if (ratedTile.paddingPenalty == bestPaddingPenalty &&
- bestRatedTile.productMxNxK < ratedTile.productMxNxK) {
- bestRatedTile = ratedTile;
- }
- }
- // Sanity check. This assert can only fail if there's a programming mistake
- // locally here.
- assert(bestRatedTile.paddingPenalty == bestPaddingPenalty);
- return bestRatedTile;
-}
-
static SmallVector<TileMxNxK>
enumerateMatmulTileMxNxK(IREE::Encoding::EncodingAttr encoding,
IREE::HAL::ExecutableTargetAttr target) {
@@ -410,9 +274,6 @@
}
// Enumerate available tile shapes for the given encoding and target.
SmallVector<Type> elementTypes = encoding.getElementTypesArray();
- if (isVMVXBackend(target)) {
- return enumerateMatmulTilesVMVX(*cDims, encoding, target);
- }
if (isAArch64(target)) {
return enumerateMatmulTileArm64(elementTypes, target);
}
@@ -442,8 +303,8 @@
auto narrowDim = IREE::Encoding::getMatmulNarrowDim(encoding);
// Choose a final matmul TileMxNxK from the above-enumarated tile shapes,
// taking narrow dimensions into account.
- TileMxNxK chosenTileMxNxK = chooseMatmulTile(enumeratedTileMxNxK, narrowDim,
- encoding.getRoundDimsToArray());
+ TileMxNxK chosenTileMxNxK = IREE::Codegen::chooseMatmulTile(
+ enumeratedTileMxNxK, narrowDim, encoding.getRoundDimsToArray());
// Map the matmul TileMxNxK to an actual tile shape for the tensor at hand,
// based on its operand index in the matmul.
@@ -481,9 +342,15 @@
// 2. We use ukernels, and this allows writing 2x fewer narrow ukernels.
// 3. Heuristics for cache-friendly dispatch tiling can get complex on CPU,
// so it is nice that they have fewer narrow cases to consider.
+ IREE::Codegen::LayoutAttrInterface layoutAttr;
+ if (isVMVXBackend(targetAttr)) {
+ layoutAttr = cast<IREE::Codegen::LayoutAttrInterface>(
+ IREE::CPU::VMVXEncodingLayoutAttr::get(ctx,
+ targetAttr.getConfiguration()));
+ }
MaterializeEncodingTypeConverter typeConverter(
materializeEncodingForTarget, targetAttr, /*transposeNarrowN=*/true,
- /*layoutAttr=*/{});
+ layoutAttr);
MaterializeEncodingConversionTarget target(*ctx);
auto materializeEncodingValueFn = getMaterializeEncodingValueFn(targetAttr);
populateMaterializeEncodingIntoPackUnPackPatterns(
@@ -547,8 +414,9 @@
: public impl::CPUMaterializeHostEncodingPassBase<
CPUMaterializeHostEncodingPass> {
void getDependentDialects(DialectRegistry ®istry) const override {
- registry.insert<arith::ArithDialect, tensor::TensorDialect,
- IREE::Codegen::IREECodegenDialect>();
+ registry
+ .insert<arith::ArithDialect, tensor::TensorDialect,
+ IREE::Codegen::IREECodegenDialect, IREE::CPU::IREECPUDialect>();
}
void runOnOperation() override {
@@ -607,8 +475,9 @@
: public impl::CPUMaterializeDeviceEncodingPassBase<
CPUMaterializeDeviceEncodingPass> {
void getDependentDialects(DialectRegistry ®istry) const override {
- registry.insert<arith::ArithDialect, tensor::TensorDialect,
- IREE::Codegen::IREECodegenDialect>();
+ registry
+ .insert<arith::ArithDialect, tensor::TensorDialect,
+ IREE::Codegen::IREECodegenDialect, IREE::CPU::IREECPUDialect>();
}
void runOnOperation() override {
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Dialect/CPU/BUILD.bazel
new file mode 100644
index 0000000..7c309b6
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/BUILD.bazel
@@ -0,0 +1,11 @@
+# Copyright 2024 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
+
+package(
+ default_visibility = ["//visibility:public"],
+ features = ["layering_check"],
+ licenses = ["notice"], # Apache 2.0
+)
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Dialect/CPU/CMakeLists.txt
new file mode 100644
index 0000000..976e774
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/CMakeLists.txt
@@ -0,0 +1,13 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from #
+# compiler/src/iree/compiler/Codegen/Dialect/CPU/BUILD.bazel #
+# #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary #
+# CMake-only content. #
+# #
+# To disable autogeneration for this file entirely, delete this header. #
+################################################################################
+
+iree_add_all_subdirs()
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/BUILD.bazel
new file mode 100644
index 0000000..676a4fd
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/BUILD.bazel
@@ -0,0 +1,105 @@
+# Copyright 2024 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
+
+load("//build_tools/bazel:build_defs.oss.bzl", "iree_compiler_cc_library", "iree_gentbl_cc_library", "iree_td_library")
+load("//build_tools/bazel:enforce_glob.bzl", "enforce_glob")
+
+package(
+ default_visibility = ["//visibility:public"],
+ features = ["layering_check"],
+ licenses = ["notice"], # Apache 2.0
+)
+
+exports_files([
+ "IREECPUAttrs.td",
+ "IREECPUDialect.td",
+])
+
+iree_td_library(
+ name = "td_files",
+ srcs = enforce_glob(
+ [
+ "IREECPUAttrs.td",
+ "IREECPUDialect.td",
+ ],
+ include = ["*.td"],
+ ),
+ deps = [
+ "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:td_files",
+ "@llvm-project//mlir:DialectUtilsTdFiles",
+ "@llvm-project//mlir:OpBaseTdFiles",
+ ],
+)
+
+iree_compiler_cc_library(
+ name = "IREECPUDialect",
+ srcs = [
+ "IREECPUAttrs.cpp",
+ "IREECPUDialect.cpp",
+ ],
+ hdrs = [
+ "IREECPUDialect.h",
+ "IREECPUTypes.h",
+ ],
+ textual_hdrs = [
+ "IREECPUAttrs.cpp.inc",
+ "IREECPUAttrs.h.inc",
+ "IREECPUDialect.cpp.inc",
+ "IREECPUDialect.h.inc",
+ ],
+ deps = [
+ ":IREECPUAttrs",
+ ":IREECPUDialectGen",
+ "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
+ "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils",
+ "//compiler/src/iree/compiler/Dialect/Encoding/IR",
+ "@llvm-project//llvm:Support",
+ "@llvm-project//mlir:AffineDialect",
+ "@llvm-project//mlir:DialectUtils",
+ "@llvm-project//mlir:IR",
+ "@llvm-project//mlir:LinalgDialect",
+ "@llvm-project//mlir:Parser",
+ "@llvm-project//mlir:Support",
+ "@llvm-project//mlir:TensorDialect",
+ ],
+)
+
+iree_gentbl_cc_library(
+ name = "IREECPUDialectGen",
+ tbl_outs = [
+ (
+ ["--gen-dialect-decls"],
+ "IREECPUDialect.h.inc",
+ ),
+ (
+ ["--gen-dialect-defs"],
+ "IREECPUDialect.cpp.inc",
+ ),
+ ],
+ tblgen = "@llvm-project//mlir:mlir-tblgen",
+ td_file = "IREECPUDialect.td",
+ deps = [":td_files"],
+)
+
+iree_gentbl_cc_library(
+ name = "IREECPUAttrs",
+ tbl_outs = [
+ (
+ ["--gen-attrdef-decls"],
+ "IREECPUAttrs.h.inc",
+ ),
+ (
+ ["--gen-attrdef-defs"],
+ "IREECPUAttrs.cpp.inc",
+ ),
+ ],
+ tblgen = "@llvm-project//mlir:mlir-tblgen",
+ td_file = "IREECPUAttrs.td",
+ deps = [
+ ":td_files",
+ "@llvm-project//mlir:SCFDeviceMappingInterfacesIncGen",
+ ],
+)
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/CMakeLists.txt
new file mode 100644
index 0000000..13ac5f0
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/CMakeLists.txt
@@ -0,0 +1,63 @@
+################################################################################
+# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from #
+# compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/BUILD.bazel #
+# #
+# Use iree_cmake_extra_content from iree/build_defs.oss.bzl to add arbitrary #
+# CMake-only content. #
+# #
+# To disable autogeneration for this file entirely, delete this header. #
+################################################################################
+
+iree_add_all_subdirs()
+
+iree_cc_library(
+ NAME
+ IREECPUDialect
+ HDRS
+ "IREECPUDialect.h"
+ "IREECPUTypes.h"
+ TEXTUAL_HDRS
+ "IREECPUAttrs.cpp.inc"
+ "IREECPUAttrs.h.inc"
+ "IREECPUDialect.cpp.inc"
+ "IREECPUDialect.h.inc"
+ SRCS
+ "IREECPUAttrs.cpp"
+ "IREECPUDialect.cpp"
+ DEPS
+ ::IREECPUAttrs
+ ::IREECPUDialectGen
+ LLVMSupport
+ MLIRAffineDialect
+ MLIRIR
+ MLIRLinalgDialect
+ MLIRParser
+ MLIRSupport
+ MLIRTensorDialect
+ iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
+ iree::compiler::Codegen::Dialect::Codegen::Utils
+ iree::compiler::Dialect::Encoding::IR
+ PUBLIC
+)
+
+iree_tablegen_library(
+ NAME
+ IREECPUDialectGen
+ TD_FILE
+ "IREECPUDialect.td"
+ OUTS
+ --gen-dialect-decls IREECPUDialect.h.inc
+ --gen-dialect-defs IREECPUDialect.cpp.inc
+)
+
+iree_tablegen_library(
+ NAME
+ IREECPUAttrs
+ TD_FILE
+ "IREECPUAttrs.td"
+ OUTS
+ --gen-attrdef-decls IREECPUAttrs.h.inc
+ --gen-attrdef-defs IREECPUAttrs.cpp.inc
+)
+
+### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp
new file mode 100644
index 0000000..2ff412b
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp
@@ -0,0 +1,35 @@
+// Copyright 2024 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/Dialect/CPU/IR/IREECPUDialect.h"
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenTypes.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h"
+#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/Support/LLVM.h"
+
+#define GET_ATTRDEF_CLASSES
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp.inc"
+
+namespace mlir::iree_compiler::IREE::CPU {
+
+//===----------------------------------------------------------------------===//
+// Attribute Registration
+//===----------------------------------------------------------------------===//
+
+void IREECPUDialect::registerAttributes() {
+ addAttributes<
+#define GET_ATTRDEF_LIST
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.cpp.inc" // IWYU pragma: keep
+ >();
+}
+
+} // namespace mlir::iree_compiler::IREE::CPU
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.td b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.td
new file mode 100644
index 0000000..a87b1e0
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.td
@@ -0,0 +1,37 @@
+// Copyright 2024 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_DIALECT_CPU_IREECPUATTRS
+#define IREE_COMPILER_CODEGEN_DIALECT_CPU_IREECPUATTRS
+
+include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.td"
+include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.td"
+include "mlir/IR/AttrTypeBase.td"
+
+//===----------------------------------------------------------------------===//
+// iree_cpu.encoding_layout_attr
+//===----------------------------------------------------------------------===//
+
+def IREECPU_VMVXEncodingLayoutAttr :
+ AttrDef<IREECPU_Dialect, "VMVXEncodingLayout"> {
+ let mnemonic = "vmvx_encoding_layout";
+ let summary = "The encoding layout attribute for VMVX backend.";
+ let description = [{
+ This attribute can implement any layout interface methods for data-tiling,
+ e.g., Codegen::LayoutAttrInterface, etc. They are implemented through
+ external model mechanism See the implementation in
+ compiler/Codegen/ExternalInterfaces/*.
+ }];
+
+ let assemblyFormat = "`<` struct(params) `>`";
+
+ let parameters = (ins
+ OptionalParameter<"DictionaryAttr", "Executable target configuration. It is "
+ "expected to be used in a pass scope, but not the final IR output.">:$configuration
+ );
+}
+
+#endif // IREE_COMPILER_CODEGEN_DIALECT_CPU_IREECPUATTRS
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.cpp b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.cpp
new file mode 100644
index 0000000..a0ced19
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.cpp
@@ -0,0 +1,15 @@
+// Copyright 2024 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/Dialect/CPU/IR/IREECPUDialect.h"
+
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.cpp.inc"
+
+namespace mlir::iree_compiler::IREE::CPU {
+
+void IREECPUDialect::initialize() { registerAttributes(); }
+
+} // namespace mlir::iree_compiler::IREE::CPU
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.h b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.h
new file mode 100644
index 0000000..785dc21
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.h
@@ -0,0 +1,18 @@
+// Copyright 2024 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_DIALECT_CPU_IREECPUDIALECT_H_
+#define IREE_COMPILER_CODEGEN_DIALECT_CPU_IREECPUDIALECT_H_
+
+#include "mlir/IR/Dialect.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Support/TypeID.h"
+
+// clang-format off: must be included after all LLVM/MLIR headers
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.h.inc" // IWYU pragma: keep
+// clang-format on
+
+#endif // IREE_COMPILER_CODEGEN_DIALECT_CPU_IREECPUDIALECT_H_
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.td b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.td
new file mode 100644
index 0000000..6c77cfb
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.td
@@ -0,0 +1,34 @@
+// Copyright 2024 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_CODEGEN_DIALECT_CPU_IREECPU_DIALECT
+#define IREE_CODEGEN_DIALECT_CPU_IREECPU_DIALECT
+
+include "mlir/IR/OpBase.td"
+include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/EnumAttr.td"
+
+def IREECPU_Dialect : Dialect {
+ let name = "iree_cpu";
+ let cppNamespace = "::mlir::iree_compiler::IREE::CPU";
+
+ let summary = [{
+ A dialect for common functionality used by CPU focused IREE code generation.
+ }];
+ let description = [{
+ This dialect provides operations and attributes to aid in code generation
+ for CPU targets. The functionality in this dialect can be hardware specific,
+ but is intended to be independent of the lowering target. Late lowerings to
+ LLVM are handled separately.
+ }];
+ let useDefaultAttributePrinterParser = 1;
+
+ let extraClassDeclaration = [{
+ void registerAttributes();
+ }];
+}
+
+#endif // IREE_CODEGEN_DIALECT_CPU_IREECPU_DIALECT
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h
new file mode 100644
index 0000000..31c5963
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h
@@ -0,0 +1,28 @@
+// Copyright 2024 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_DIALECT_CPU_IREECPUTYPES_H_
+#define IREE_COMPILER_CODEGEN_DIALECT_CPU_IREECPUTYPES_H_
+
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h"
+
+// clang-format off
+#define GET_ATTRDEF_CLASSES
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUAttrs.h.inc"
+#undef GET_ATTRDEF_CLASSES
+// clang-format on
+
+namespace mlir::iree_compiler::IREE::CPU {
+
+//===----------------------------------------------------------------------===//
+// Utilities.
+//===----------------------------------------------------------------------===//
+
+} // namespace mlir::iree_compiler::IREE::CPU
+
+#endif // IREE_COMPILER_CODEGEN_DIALECT_CPU_IREECPUTYPES_H_
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.cpp b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.cpp
index 7976b7e..626280e 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.cpp
@@ -6,6 +6,7 @@
#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h"
#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Debug.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
@@ -15,6 +16,8 @@
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
+#define DEBUG_TYPE "iree-codegen-dialect-codegen-utils"
+
namespace mlir::iree_compiler::IREE::Codegen {
//===----------------------------------------------------------------------===//
@@ -404,6 +407,106 @@
return expandedValue;
}
+TileMxNxK chooseMatmulTile(ArrayRef<TileMxNxK> enumeratedTiles,
+ IREE::Encoding::MatmulNarrowDim narrowDim,
+ ArrayRef<int64_t> hostDefinedUpperBound) {
+ assert((hostDefinedUpperBound.empty() || hostDefinedUpperBound.size() >= 3) &&
+ "expected hostDefinedUpperBound is empty or has upper bound for {M, "
+ "N, K}");
+ // Handle narrow-N by transposing to reduce to narrow-M. Note: the
+ // enumeratedTiles currently only enumerate narrow-M cases.
+ if (narrowDim.isN()) {
+ SmallVector<int64_t> newHostDefinedUpperBound(hostDefinedUpperBound);
+ std::swap(newHostDefinedUpperBound[0], newHostDefinedUpperBound[1]);
+ narrowDim.dim = IREE::Encoding::MatmulNarrowDim::Dim::M;
+ TileMxNxK tile =
+ chooseMatmulTile(enumeratedTiles, narrowDim, newHostDefinedUpperBound);
+ std::swap(tile.M, tile.N);
+ return tile;
+ }
+ // Handle kDynamic: currently this is only used with VMVX, where there is only
+ // one enumerated tile and it has all three M/N/K dimensions dynamic, so for
+ // now we only support that. Generalize that as needed when more dynamic tile
+ // sizes are used outside of VMVX, e.g. perhaps some day with Arm SVE. Decide
+ // how to incorporate the handling of kDynamic in the cost-model evaluation
+ // below to decide when to prefer a dynamic vs a static tile shape.
+ for (auto tile : enumeratedTiles) {
+ if (ShapedType::isDynamic(tile.M) || ShapedType::isDynamic(tile.N) ||
+ ShapedType::isDynamic(tile.K)) {
+ assert(enumeratedTiles.size() == 1);
+ assert(ShapedType::isDynamic(tile.M) && ShapedType::isDynamic(tile.N) &&
+ ShapedType::isDynamic(tile.K));
+ return tile;
+ }
+ }
+ // We're going to "rate" the enumerated tiles.
+ struct RatedTileMxNxK : TileMxNxK {
+ RatedTileMxNxK() {}
+ RatedTileMxNxK(TileMxNxK tile) : TileMxNxK(tile) {}
+ // Penalize tiles that are wider in the M dimension than matmulNarrowM.
+ int64_t paddingPenalty = 0;
+ // Favor larger tiles, as long as they still minimize paddingPenalty.
+ int64_t productMxNxK = 0;
+ };
+ SmallVector<RatedTileMxNxK> ratedTiles;
+ ratedTiles.reserve(enumeratedTiles.size());
+ int64_t bestPaddingPenalty = INT64_MAX;
+ int64_t mUB = INT64_MAX;
+ int64_t nUB = INT64_MAX;
+ int64_t kUB = INT64_MAX;
+ if (!hostDefinedUpperBound.empty()) {
+ mUB = hostDefinedUpperBound[0];
+ nUB = hostDefinedUpperBound[1];
+ kUB = hostDefinedUpperBound[2];
+ }
+ for (auto tile : enumeratedTiles) {
+ if (tile.M > mUB || tile.N > nUB || tile.K > kUB) {
+ LLVM_DEBUG(llvm::dbgs() << "[" << DEBUG_TYPE << "]: tile (";
+ llvm::interleaveComma(
+ ArrayRef<int64_t>{tile.M, tile.N, tile.K}, llvm::dbgs());
+ llvm::dbgs()
+ << ") is skipped because it is not valid for upper_bound (";
+ llvm::interleaveComma(ArrayRef<int64_t>{mUB, nUB, kUB},
+ llvm::dbgs());
+ llvm::dbgs() << ")\n");
+ continue;
+ }
+ RatedTileMxNxK ratedTile(tile);
+ ratedTile.paddingPenalty = 0;
+ // If we are choosing a tile for a narrow-M case, we want to minimize
+ // padding along the M dimension.
+ // The PowerOf2Ceil is so that we are OK with padding up to the next
+ // power of two, we just try to avoid padding beyond that. For example,
+ // if matmulNarrowM==7 and we have enumerated tiles with M=8,4,2,1, we
+ // are OK with the tile that has M==8 even though it requires some padding.
+ // Otherwise, we would be penalizing the tiles with M==8,4,2 and we would
+ // end up selecting the vecmat tile (M==1) for that case!
+ if (narrowDim) {
+ ratedTile.paddingPenalty =
+ std::max<int64_t>(tile.M - llvm::PowerOf2Ceil(narrowDim.size), 0);
+ }
+ ratedTile.productMxNxK = tile.M * tile.N * tile.K;
+ ratedTiles.push_back(ratedTile);
+ LLVM_DEBUG(llvm::dbgs() << "candidate: "; llvm::interleaveComma(
+ ArrayRef<int64_t>{tile.M, tile.N, tile.K}, llvm::dbgs());
+ llvm::dbgs() << " penalty:" << ratedTile.paddingPenalty << "\n");
+ bestPaddingPenalty = std::min(bestPaddingPenalty, ratedTile.paddingPenalty);
+ }
+ RatedTileMxNxK bestRatedTile;
+ for (auto ratedTile : ratedTiles) {
+ // Choose only among tiles that minimize paddingPenalty. Among those,
+ // maximize productMxNxK.
+ if (ratedTile.paddingPenalty == bestPaddingPenalty &&
+ bestRatedTile.productMxNxK < ratedTile.productMxNxK) {
+ bestRatedTile = ratedTile;
+ }
+ }
+ // Sanity check. This assert can only fail if there's a programming mistake
+ // locally here.
+ assert(bestRatedTile.paddingPenalty == bestPaddingPenalty);
+ return bestRatedTile;
+}
+
FailureOr<Operation *>
lowerContractionOpWithEncoding(OpBuilder &builder, linalg::LinalgOp linalgOp,
ValueRange operands, bool transposeNarrowN,
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h
index 8498a95..f08aec2 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h
@@ -78,6 +78,20 @@
// Operation Lowering Utilities.
//===----------------------------------------------------------------------===//
+// TODO(hanchung): The below methods are exposed to public because they are
+// shared between MaterializeEncodingIntoPackUnPack.cpp.cpp and
+// CPUEncodingExternalModels.cpp. They will be moved to other places after all
+// the CPU backends implement their layout attributes.
+
+/// Returns the best TileMxNxK from `enumeratedTiles` pool. If the
+/// `hostDefinedUpperBound` is not empty, the chosen tile sizes can not be
+/// greater than the values.
+/// TODO(#16933): Remove `hostDefinedUpperBound` once we can propagate such
+/// information to host. For now, they are defined by host.
+TileMxNxK chooseMatmulTile(ArrayRef<TileMxNxK> enumeratedTiles,
+ IREE::Encoding::MatmulNarrowDim narrowDim,
+ ArrayRef<int64_t> hostDefinedUpperBound = {});
+
FailureOr<Operation *>
lowerContractionOpWithEncoding(OpBuilder &builder, linalg::LinalgOp linalgOp,
ValueRange operands, bool transposeNarrowN,
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/BUILD.bazel b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/BUILD.bazel
index 9b47544..640e1ef 100644
--- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/BUILD.bazel
@@ -15,17 +15,21 @@
iree_compiler_cc_library(
name = "ExternalModels",
srcs = [
+ "CPUEncodingExternalModels.cpp",
"GPUEncodingExternalModels.cpp",
"Interfaces.cpp",
],
hdrs = [
+ "CPUEncodingExternalModels.h",
"GPUEncodingExternalModels.h",
"Interfaces.h",
],
deps = [
+ "//compiler/src/iree/compiler/Codegen/Dialect/CPU/IR:IREECPUDialect",
"//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
"//compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils",
"//compiler/src/iree/compiler/Codegen/Dialect/GPU/IR:IREEGPUDialect",
+ "//compiler/src/iree/compiler/Codegen/Utils",
"//compiler/src/iree/compiler/Dialect/Encoding/IR",
"@llvm-project//llvm:Support",
"@llvm-project//mlir:IR",
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CMakeLists.txt
index c637517..46d0627 100644
--- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CMakeLists.txt
@@ -14,9 +14,11 @@
NAME
ExternalModels
HDRS
+ "CPUEncodingExternalModels.h"
"GPUEncodingExternalModels.h"
"Interfaces.h"
SRCS
+ "CPUEncodingExternalModels.cpp"
"GPUEncodingExternalModels.cpp"
"Interfaces.cpp"
DEPS
@@ -24,9 +26,11 @@
MLIRIR
MLIRLinalgDialect
MLIRTensorDialect
+ iree::compiler::Codegen::Dialect::CPU::IR::IREECPUDialect
iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
iree::compiler::Codegen::Dialect::Codegen::Utils
iree::compiler::Codegen::Dialect::GPU::IR::IREEGPUDialect
+ iree::compiler::Codegen::Utils
iree::compiler::Dialect::Encoding::IR
PUBLIC
)
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
new file mode 100644
index 0000000..1e09bc6
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
@@ -0,0 +1,123 @@
+// Copyright 2024 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/CPUEncodingExternalModels.h"
+
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenTypes.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h"
+#include "iree/compiler/Codegen/Utils/Utils.h"
+#include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
+#include "llvm/Support/Debug.h"
+#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
+
+#define DEBUG_TYPE "iree-gpu-encoding-external-models"
+
+namespace mlir::iree_compiler::IREE::CPU {
+
+using Codegen::MaterializeEncodingInfo;
+using Codegen::TileMxNxK;
+
+namespace {
+
+//===----------------------------------------------------------------------===//
+// Interface methods implementaion for iree_cpu.vmvx_encoding_layout.
+//===----------------------------------------------------------------------===//
+
+// Enumerate tile sizes to choose from when no specific architecture is
+// targeted. For narrow-{M,N} cases, this only enumerates on narrow M. The
+// narrow-N cases are handled by transposition in chooseMatmulTile.
+static SmallVector<TileMxNxK>
+enumerateVMVXMatmulTiles(linalg::ContractionDimensions cDims,
+ IREE::Encoding::EncodingAttr encoding,
+ DictionaryAttr config) {
+ bool hasUkernelSupport = hasUkernel(config);
+
+ // TODO(hanchung): The ukernel path does not support 3d
+ // codegen.query_tile_sizes op, so we disable dynamic tile shapes for
+ // batch_matmul. Also, they are not set up for narrow M/N matmul, so it is
+ // disabled when it is the case.
+ if (!cDims.batch.empty() || getMatmulNarrowDim(encoding)) {
+ hasUkernelSupport = false;
+ }
+ if (hasUkernelSupport) {
+ // VMVX+ukernel uses dynamic tile shapes.
+ return {TileMxNxK{ShapedType::kDynamic, ShapedType::kDynamic,
+ ShapedType::kDynamic}};
+ }
+
+ return {
+ TileMxNxK{8, 8, 4}, // Some vaguely reasonable tile shape.
+ TileMxNxK{4, 8, 4}, // Truncation of the above.
+ TileMxNxK{2, 8, 4}, // Truncation of the above.
+ TileMxNxK{1, 8, 4}, // Truncation of the above.
+ };
+}
+
+struct VMVXDeviceEncodingLayoutAttrInterface
+ : public Codegen::LayoutAttrInterface::ExternalModel<
+ VMVXDeviceEncodingLayoutAttrInterface, VMVXEncodingLayoutAttr> {
+ MaterializeEncodingInfo getEncodingInfo(Attribute attr,
+ RankedTensorType type) const {
+ auto layoutAttr = cast<VMVXEncodingLayoutAttr>(attr);
+ auto encoding = llvm::dyn_cast_or_null<IREE::Encoding::EncodingAttr>(
+ type.getEncoding());
+
+ MaterializeEncodingInfo info;
+ if (!encoding) {
+ return info;
+ }
+
+ // We only know about contractions with {Batch, M, N, K} <= 1 at the moment.
+ auto cDims = getEncodingContractionDims(encoding);
+ if (failed(cDims) || cDims->batch.size() > 1 || cDims->m.size() > 1 ||
+ cDims->n.size() > 1 || cDims->k.size() > 1) {
+ return info;
+ }
+
+ SmallVector<TileMxNxK> enumeratedTileMxNxK = enumerateVMVXMatmulTiles(
+ cDims.value(), encoding, layoutAttr.getConfiguration());
+ if (enumeratedTileMxNxK.empty()) {
+ return info;
+ }
+ auto narrowDim = IREE::Encoding::getMatmulNarrowDim(encoding);
+ // Choose a final matmul TileMxNxK from the above-enumarated tile shapes,
+ // taking narrow dimensions into account.
+ TileMxNxK chosenTileMxNxK = chooseMatmulTile(
+ enumeratedTileMxNxK, narrowDim, encoding.getRoundDimsToArray());
+ return getEncodingInfoForMatmul(encoding, chosenTileMxNxK);
+ }
+
+ Operation *lowerOp(Attribute attr, OpBuilder &b, Operation *op,
+ TypeRange convertedResTypes,
+ ValueRange convertedOperands) const {
+ auto layoutAttr = cast<VMVXEncodingLayoutAttr>(attr);
+ auto linalgOp = llvm::dyn_cast<linalg::LinalgOp>(op);
+ if (!linalgOp) {
+ return nullptr;
+ }
+
+ auto resolver =
+ [&](RankedTensorType type) -> FailureOr<MaterializeEncodingInfo> {
+ return getEncodingInfo(layoutAttr, type);
+ };
+ FailureOr<Operation *> newOp = Codegen::lowerContractionOpWithEncoding(
+ b, linalgOp, convertedOperands, /*transposeNarrowN=*/true, resolver);
+ return newOp.value_or(nullptr);
+ }
+};
+
+} // namespace
+
+void registerCPUEncodingExternalModels(DialectRegistry ®istry) {
+ registry.addExtension(
+ +[](MLIRContext *ctx, IREE::CPU::IREECPUDialect *dialect) {
+ IREE::CPU::VMVXEncodingLayoutAttr::attachInterface<
+ VMVXDeviceEncodingLayoutAttrInterface>(*ctx);
+ });
+}
+
+} // namespace mlir::iree_compiler::IREE::CPU
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.h b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.h
new file mode 100644
index 0000000..715ed3f
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.h
@@ -0,0 +1,20 @@
+// Copyright 2024 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_CPUENCODINGEXTERNALMODELS_H_
+#define IREE_COMPILER_CODEGEN_EXTERNALINTERFACES_CPUENCODINGEXTERNALMODELS_H_
+
+namespace mlir {
+class DialectRegistry;
+} // namespace mlir
+
+namespace mlir::iree_compiler::IREE::CPU {
+
+void registerCPUEncodingExternalModels(DialectRegistry ®istry);
+
+} // namespace mlir::iree_compiler::IREE::CPU
+
+#endif // IREE_COMPILER_CODEGEN_EXTERNALINTERFACES_CPUENCODINGEXTERNALMODELS_H_
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Interfaces.cpp b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Interfaces.cpp
index 2d19708..2f3b363 100644
--- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Interfaces.cpp
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Interfaces.cpp
@@ -6,12 +6,14 @@
#include "iree/compiler/Codegen/ExternalInterfaces/Interfaces.h"
+#include "iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.h"
#include "iree/compiler/Codegen/ExternalInterfaces/GPUEncodingExternalModels.h"
namespace mlir::iree_compiler {
void registerCodegenExternalInterfaces(DialectRegistry ®istry) {
IREE::GPU::registerGPUEncodingExternalModels(registry);
+ IREE::CPU::registerCPUEncodingExternalModels(registry);
}
} // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp b/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp
index f86f447..4145311 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp
@@ -11,10 +11,12 @@
#include "iree/compiler/Codegen/Interfaces/UKernelOpInterface.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
+#include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtDialect.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Casting.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
@@ -28,6 +30,8 @@
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/AffineExprVisitor.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Interfaces/TilingInterface.h"
@@ -61,53 +65,77 @@
return func.isPublic() && getEntryPoint(func);
}
-std::optional<StringAttr>
-getConfigStringAttr(IREE::HAL::ExecutableTargetAttr targetAttr,
- StringRef stringAttr) {
- if (!targetAttr)
+std::optional<StringAttr> getConfigStringAttr(Attribute srcAttr,
+ StringRef stringAttr) {
+ if (!srcAttr) {
return std::nullopt;
- auto config = targetAttr.getConfiguration();
- if (!config)
+ }
+ auto targetAttr = dyn_cast<IREE::HAL::ExecutableTargetAttr>(srcAttr);
+ DictionaryAttr config;
+ if (targetAttr) {
+ config = targetAttr.getConfiguration();
+ } else {
+ config = dyn_cast<DictionaryAttr>(srcAttr);
+ }
+ if (!config) {
return std::nullopt;
+ }
auto attr = config.getAs<StringAttr>(stringAttr);
- if (!attr)
+ if (!attr) {
return std::nullopt;
+ }
return attr;
}
-std::optional<IntegerAttr>
-getConfigIntegerAttr(IREE::HAL::ExecutableTargetAttr targetAttr,
- StringRef integerAttr) {
- if (!targetAttr)
+std::optional<IntegerAttr> getConfigIntegerAttr(Attribute srcAttr,
+ StringRef integerAttr) {
+ if (!srcAttr) {
return std::nullopt;
- auto config = targetAttr.getConfiguration();
- if (!config)
+ }
+ auto targetAttr = dyn_cast<IREE::HAL::ExecutableTargetAttr>(srcAttr);
+ DictionaryAttr config;
+ if (targetAttr) {
+ config = targetAttr.getConfiguration();
+ } else {
+ config = dyn_cast<DictionaryAttr>(srcAttr);
+ }
+ if (!config) {
return std::nullopt;
+ }
auto attr = config.getAs<IntegerAttr>(integerAttr);
- if (!attr)
+ if (!attr) {
return std::nullopt;
+ }
return attr;
}
-std::optional<BoolAttr>
-getConfigBoolAttr(IREE::HAL::ExecutableTargetAttr targetAttr,
- StringRef integerAttr) {
- if (!targetAttr)
+std::optional<BoolAttr> getConfigBoolAttr(Attribute srcAttr,
+ StringRef boolAttr) {
+ if (!srcAttr) {
return std::nullopt;
- auto config = targetAttr.getConfiguration();
- if (!config)
+ }
+ auto targetAttr = dyn_cast<IREE::HAL::ExecutableTargetAttr>(srcAttr);
+ DictionaryAttr config;
+ if (targetAttr) {
+ config = targetAttr.getConfiguration();
+ } else {
+ config = dyn_cast<DictionaryAttr>(srcAttr);
+ }
+ if (!config) {
return std::nullopt;
- auto attr = config.getAs<BoolAttr>(integerAttr);
- if (!attr)
+ }
+ auto attr = config.getAs<BoolAttr>(boolAttr);
+ if (!attr) {
return std::nullopt;
+ }
return attr;
}
-std::optional<llvm::Triple>
-getTargetTriple(IREE::HAL::ExecutableTargetAttr targetAttr) {
- auto triple = getConfigStringAttr(targetAttr, "target_triple");
- if (!triple)
+std::optional<llvm::Triple> getTargetTriple(Attribute attr) {
+ auto triple = getConfigStringAttr(attr, "target_triple");
+ if (!triple) {
return std::nullopt;
+ }
return llvm::Triple(triple.value().str());
}
@@ -141,24 +169,30 @@
return targetAttr && targetAttr.getBackend().getValue().starts_with("rocm");
}
-static const char *
-getDefaultEnabledUkernels(IREE::HAL::ExecutableTargetAttr targetAttr) {
+static const char *getDefaultEnabledUkernels(Attribute attr) {
+ const char *kNone = "none";
+ if (!attr) {
+ return kNone;
+ }
+ auto targetAttr = dyn_cast<IREE::HAL::ExecutableTargetAttr>(attr);
+ if (!targetAttr) {
+ return kNone;
+ }
if (isX86_64(targetAttr)) {
return "mmt4d";
}
if (isAArch64(targetAttr)) {
if (hasFeature(targetAttr, "+sve") || hasFeature(targetAttr, "+sve2") ||
hasFeature(targetAttr, "+sme")) {
- return "none";
+ return kNone;
}
return "mmt4d";
}
- return "none";
+ return kNone;
}
-bool hasUkernel(IREE::HAL::ExecutableTargetAttr targetAttr,
- StringRef ukernelName) {
- auto enabledUkernels = getConfigStringAttr(targetAttr, "ukernels");
+bool hasUkernel(Attribute attr, StringRef ukernelName) {
+ auto enabledUkernels = getConfigStringAttr(attr, "ukernels");
StringRef enabledUkernelsStr;
if (enabledUkernels) {
enabledUkernelsStr = enabledUkernels->getValue();
@@ -167,7 +201,7 @@
}
// Resolve `default`.
if (enabledUkernelsStr == "default") {
- enabledUkernelsStr = getDefaultEnabledUkernels(targetAttr);
+ enabledUkernelsStr = getDefaultEnabledUkernels(attr);
}
// Resolve `none`.
if (enabledUkernelsStr == "none") {
diff --git a/compiler/src/iree/compiler/Codegen/Utils/Utils.h b/compiler/src/iree/compiler/Codegen/Utils/Utils.h
index 4603429..6a43cb4 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/Utils.h
+++ b/compiler/src/iree/compiler/Codegen/Utils/Utils.h
@@ -37,27 +37,24 @@
std::optional<IREE::HAL::ExecutableExportOp>
getEntryPoint(mlir::FunctionOpInterface funcOp);
-/// Returns the StringAttr with the name `stringAttr` in the `targetAttr`, if
-/// found.
-std::optional<StringAttr>
-getConfigStringAttr(IREE::HAL::ExecutableTargetAttr targetAttr,
- StringRef stringAttr);
+/// Returns the StringAttr with the name `stringAttr` in the `srcAttr`, if
+/// found. The `srcAttr` can be either IREE::HAL::ExecutableTargetAttr or
+/// DictionaryAttr.
+std::optional<StringAttr> getConfigStringAttr(Attribute srcAttr,
+ StringRef stringAttr);
-/// Returns the IntegerAttr with the name `integerAttr` in the `targetAttr`, if
+/// Returns the IntegerAttr with the name `integerAttr` in the `srcAttr`, if
/// found.
-std::optional<IntegerAttr>
-getConfigIntegerAttr(IREE::HAL::ExecutableTargetAttr targetAttr,
- StringRef integerAttr);
+std::optional<IntegerAttr> getConfigIntegerAttr(Attribute srcAttr,
+ StringRef integerAttr);
-/// Returns the BoolAttr with the name `integerAttr` in the `targetAttr`, if
+/// Returns the BoolAttr with the name `boolAttr` in the `srcAttr`, if
/// found.
-std::optional<BoolAttr>
-getConfigBoolAttr(IREE::HAL::ExecutableTargetAttr targetAttr,
- StringRef integerAttr);
+std::optional<BoolAttr> getConfigBoolAttr(Attribute srcAttr,
+ StringRef boolAttr);
-/// Returns the LLVM Target triple associated with the `targetAttr`, if set.
-std::optional<llvm::Triple>
-getTargetTriple(IREE::HAL::ExecutableTargetAttr targetAttr);
+/// Returns the LLVM Target triple associated with the `attr`, if set.
+std::optional<llvm::Triple> getTargetTriple(Attribute attr);
/// Returns the target architecture name, in IREE_ARCH convention, from the
/// given target triple.
@@ -72,8 +69,7 @@
// Returns true if the ukernel with given `ukernelName` is enabled.
// If `ukernelName` is empty (the default), returns true if any ukernel
// is enabled at all.
-bool hasUkernel(IREE::HAL::ExecutableTargetAttr targetAttr,
- StringRef ukernelName = "");
+bool hasUkernel(Attribute attr, StringRef ukernelName = "");
/// Returns the CPU target features associated with the `targetAttr`, if set.
std::optional<StringRef>
diff --git a/compiler/src/iree/compiler/Tools/BUILD.bazel b/compiler/src/iree/compiler/Tools/BUILD.bazel
index 7a813c2..58ac050 100644
--- a/compiler/src/iree/compiler/Tools/BUILD.bazel
+++ b/compiler/src/iree/compiler/Tools/BUILD.bazel
@@ -34,6 +34,7 @@
deps = [
"//compiler/src/iree/compiler/Bindings/Native/Transforms",
"//compiler/src/iree/compiler/Bindings/TFLite/Transforms",
+ "//compiler/src/iree/compiler/Codegen/Dialect/CPU/IR:IREECPUDialect",
"//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
"//compiler/src/iree/compiler/Codegen/Dialect/GPU/IR:IREEGPUDialect",
"//compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR:IREEVectorExtDialect",
diff --git a/compiler/src/iree/compiler/Tools/init_iree_dialects.h b/compiler/src/iree/compiler/Tools/init_iree_dialects.h
index 0bc4123..0312382 100644
--- a/compiler/src/iree/compiler/Tools/init_iree_dialects.h
+++ b/compiler/src/iree/compiler/Tools/init_iree_dialects.h
@@ -14,6 +14,7 @@
#include "iree-dialects/Dialect/Input/InputDialect.h"
#include "iree-dialects/Dialect/LinalgTransform/Passes.h"
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUDialect.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
#include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
@@ -39,7 +40,8 @@
// Add all the IREE dialects to the provided registry.
inline void registerIreeDialects(DialectRegistry ®istry) {
// clang-format off
- registry.insert<IREE::Codegen::IREECodegenDialect,
+ registry.insert<IREE::CPU::IREECPUDialect,
+ IREE::Codegen::IREECodegenDialect,
IREE::Flow::FlowDialect,
IREE::GPU::IREEGPUDialect,
IREE::HAL::HALDialect,
