[DT][Encoding] Use layouts to calculate storage size when it is present. (#19686)
The revision implements the `calculateStorageSizeInBytes` method for
backend attributes, so the EncodingAttr can query the storage size from
each layout when layouts are set. It enables the storage size request
propagation from devices to hosts through interfaces.
The tests are adapted from existing tests with resolved layouts.
Additionally, it adds a new test that has two different layouts. It
computes both of them and select the maximum size.
The revision also implements the output streaming and log the final tile
size decision in CPUEncodingExternalModels, which helps debug easier.
---------
Signed-off-by: hanhanW <hanhan0912@gmail.com>
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 c515766..3e0669c 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.cpp
@@ -80,6 +80,21 @@
return !(lhs == rhs);
}
+llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
+ const MaterializeEncodingInfo &encodingInfo) {
+ os << "{innerDimsPos = [";
+ llvm::interleaveComma(encodingInfo.innerDimsPos, os);
+ os << "], innerTileSizes = [";
+ llvm::interleaveComma(encodingInfo.innerTileSizes, os);
+ os << "], outerDimsPerm = [";
+ llvm::interleaveComma(encodingInfo.outerDimsPerm, os);
+ if (encodingInfo.swizzle) {
+ os << "], swizzle = " << encodingInfo.swizzle.value();
+ }
+ os << "]}";
+ return os;
+}
+
//===----------------------------------------------------------------------===//
// Layout Utilities.
//===----------------------------------------------------------------------===//
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 8a8c309..bdb7ebb 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h
+++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h
@@ -38,6 +38,8 @@
const MaterializeEncodingInfo &rhs);
bool operator!=(const MaterializeEncodingInfo &lhs,
const MaterializeEncodingInfo &rhs);
+llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
+ const MaterializeEncodingInfo &encodingInfo);
//===----------------------------------------------------------------------===//
// Layout Utilities.
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/BUILD.bazel b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/BUILD.bazel
index ec9b711..326c97a 100644
--- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/BUILD.bazel
@@ -34,6 +34,7 @@
"//compiler/src/iree/compiler/Codegen/Utils",
"//compiler/src/iree/compiler/Dialect/Encoding/IR",
"@llvm-project//llvm:Support",
+ "@llvm-project//mlir:ArithDialect",
"@llvm-project//mlir:IR",
"@llvm-project//mlir:LinalgDialect",
"@llvm-project//mlir:TensorDialect",
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CMakeLists.txt
index 3f14e77..181b029 100644
--- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CMakeLists.txt
@@ -25,6 +25,7 @@
"Utils.cpp"
DEPS
LLVMSupport
+ MLIRArithDialect
MLIRIR
MLIRLinalgDialect
MLIRTensorDialect
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
index 187cef7..b279db3 100644
--- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp
@@ -40,6 +40,7 @@
#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
#include "llvm/Support/Debug.h"
#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
+#include "mlir/IR/BuiltinAttributes.h"
#define DEBUG_TYPE "iree-cpu-encoding-external-models"
@@ -285,6 +286,11 @@
// Sanity check. This assert can only fail if there's a programming mistake
// locally here.
assert(bestRatedTile.paddingPenalty == bestPaddingPenalty);
+ LLVM_DEBUG(
+ llvm::dbgs() << "bestRatedTile: "; llvm::interleaveComma(
+ ArrayRef<int64_t>{bestRatedTile.M, bestRatedTile.N, bestRatedTile.K},
+ llvm::dbgs());
+ llvm::dbgs() << " penalty:" << bestRatedTile.paddingPenalty << "\n");
return bestRatedTile;
}
@@ -604,6 +610,18 @@
MaterializeEncodingInfo getEncodingInfo(Attribute attr,
RankedTensorType type) const {
auto layoutAttr = cast<CPUEncodingLayoutAttr>(attr);
+
+ // If the layout is already resolved, use it directly.
+ if (auto config = layoutAttr.getConfiguration()) {
+ if (auto namedAttr = config.getNamed(kEncodingInfoAttrName)) {
+ std::optional<MaterializeEncodingInfo> info =
+ Codegen::deserializeEncodingInfo(
+ cast<DictionaryAttr>(namedAttr->getValue()));
+ assert(info && "encoding_info is invalid");
+ return info.value();
+ }
+ }
+
auto encoding = llvm::dyn_cast_or_null<IREE::Encoding::EncodingAttr>(
type.getEncoding());
@@ -655,6 +673,14 @@
struct CPUHostEncodingLayoutAttrInterface
: public IREE::Encoding::EncodingLayoutAttrInterface::ExternalModel<
CPUHostEncodingLayoutAttrInterface, CPUEncodingLayoutAttr> {
+
+ Value calculateStorageSizeInBytes(Attribute attr, Location loc,
+ OpBuilder &builder, RankedTensorType type,
+ ValueRange dynamicDims) const {
+ return calculateStorageSizeInBytesImpl(attr, loc, builder, type,
+ dynamicDims);
+ }
+
Attribute cloneWithSimplifiedConfig(Attribute attr,
DictionaryAttr config) const {
MLIRContext *ctx = attr.getContext();
@@ -711,6 +737,18 @@
MaterializeEncodingInfo getEncodingInfo(Attribute attr,
RankedTensorType type) const {
auto layoutAttr = cast<VMVXEncodingLayoutAttr>(attr);
+
+ // If the layout is already resolved, use it directly.
+ if (auto config = layoutAttr.getConfiguration()) {
+ if (auto namedAttr = config.getNamed(kEncodingInfoAttrName)) {
+ std::optional<MaterializeEncodingInfo> info =
+ Codegen::deserializeEncodingInfo(
+ cast<DictionaryAttr>(namedAttr->getValue()));
+ assert(info && "encoding_info is invalid");
+ return info.value();
+ }
+ }
+
auto encoding = llvm::dyn_cast_or_null<IREE::Encoding::EncodingAttr>(
type.getEncoding());
@@ -762,6 +800,13 @@
struct VMVXHostEncodingLayoutAttrInterface
: public IREE::Encoding::EncodingLayoutAttrInterface::ExternalModel<
VMVXHostEncodingLayoutAttrInterface, VMVXEncodingLayoutAttr> {
+ Value calculateStorageSizeInBytes(Attribute attr, Location loc,
+ OpBuilder &builder, RankedTensorType type,
+ ValueRange dynamicDims) const {
+ return calculateStorageSizeInBytesImpl(attr, loc, builder, type,
+ dynamicDims);
+ }
+
Attribute cloneWithSimplifiedConfig(Attribute attr,
DictionaryAttr config) const {
MLIRContext *ctx = attr.getContext();
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Utils.cpp b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Utils.cpp
index 12542e2..b2629aa 100644
--- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Utils.cpp
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Utils.cpp
@@ -6,20 +6,97 @@
#include "iree/compiler/Codegen/ExternalInterfaces/Utils.h"
+#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h"
#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/Dialect/Encoding/IR/EncodingTypes.h"
#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/MLIRContext.h"
namespace mlir::iree_compiler::IREE {
using Codegen::MaterializeEncodingInfo;
+Value calculateStorageSizeInBytesImpl(Attribute attr, Location loc,
+ OpBuilder &builder, RankedTensorType type,
+ ValueRange dynamicDims) {
+ auto deviceLayoutAttr = cast<IREE::Codegen::LayoutAttrInterface>(attr);
+ MaterializeEncodingInfo encodingInfo = deviceLayoutAttr.getEncodingInfo(type);
+ SmallVector<int64_t> paddedShape(type.getShape());
+ SmallVector<Value> paddedDynamicDims(dynamicDims.begin(), dynamicDims.end());
+ for (auto [dim, size] : llvm::zip_equal(encodingInfo.innerDimsPos,
+ encodingInfo.innerTileSizes)) {
+ // Only VMVX backend has dynamic inner tile sizes when ukernel is enabled.
+ // It assumes that the padding size is 16. Ideally, the logic should be
+ // moved to VMVX implementation details. However, we cook the logic here to
+ // reduce code duplication.
+ if (ShapedType::isDynamic(size)) {
+ assert(isa<IREE::CPU::VMVXEncodingLayoutAttr>(attr) &&
+ "only VMVX backend attribute can handle dynamic tile sizes");
+ size = 16;
+ }
+
+ // Do not create additional operations in the first place if the padding is
+ // not needed.
+ if (size == 1) {
+ continue;
+ }
+
+ if (type.isDynamicDim(dim)) {
+ dim = type.getDynamicDimIndex(dim);
+ auto alignment = builder.create<arith::ConstantIndexOp>(loc, size);
+ paddedDynamicDims[dim] = builder.create<arith::CeilDivSIOp>(
+ loc, paddedDynamicDims[dim], alignment);
+ paddedDynamicDims[dim] =
+ builder.create<arith::MulIOp>(loc, paddedDynamicDims[dim], alignment);
+ } else {
+ paddedShape[dim] = llvm::alignTo(paddedShape[dim], size);
+ }
+ }
+
+ constexpr int64_t kNumBitsInByte = 8;
+ int64_t numBytesPerElem = 1;
+ int64_t elementBits = type.getElementTypeBitWidth();
+ if (elementBits > kNumBitsInByte) {
+ numBytesPerElem *= elementBits / kNumBitsInByte;
+ }
+
+ int64_t staticCount = numBytesPerElem;
+ for (unsigned i = 0, e = type.getRank(); i < e; ++i) {
+ if (!type.isDynamicDim(i)) {
+ staticCount *= paddedShape[i];
+ }
+ }
+
+ Value result =
+ builder.create<arith::ConstantIndexOp>(loc, staticCount).getResult();
+ for (auto dim : paddedDynamicDims) {
+ result = builder.create<arith::MulIOp>(loc, result, dim);
+ }
+
+ // Always pack the elements back-to-back for subtypes.
+ if (elementBits < kNumBitsInByte) {
+ if (kNumBitsInByte % elementBits) {
+ assert(false && "unsupported subtype");
+ return Value();
+ }
+ Value divisor = builder.create<arith::ConstantIndexOp>(
+ loc, kNumBitsInByte / elementBits);
+ result = builder.create<arith::CeilDivUIOp>(loc, result, divisor);
+ }
+
+ return result;
+}
+
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");
+ auto strAttr = StringAttr::get(ctx, kEncodingInfoAttrName);
Attribute encodingInfoAttr =
IREE::Codegen::serializeEncodingInfo(attr.getContext(), info);
return DictionaryAttr::get(ctx, {NamedAttribute(strAttr, encodingInfoAttr)});
diff --git a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Utils.h b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Utils.h
index f4649d6..d6f3f7f 100644
--- a/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Utils.h
+++ b/compiler/src/iree/compiler/Codegen/ExternalInterfaces/Utils.h
@@ -7,11 +7,19 @@
#ifndef IREE_COMPILER_CODEGEN_EXTERNALINTERFACES_UTILS_H_
#define IREE_COMPILER_CODEGEN_EXTERNALINTERFACES_UTILS_H_
+#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/ValueRange.h"
namespace mlir::iree_compiler::IREE {
+static const char kEncodingInfoAttrName[] = "encoding_info";
+
+Value calculateStorageSizeInBytesImpl(Attribute attr, Location loc,
+ OpBuilder &builder, RankedTensorType type,
+ ValueRange dynamicDims);
+
/// Returns a dictionary attribute that contains the materialized encoding info,
/// i.e., serialized MaterializeEncodingInfo struct.
/// Requirement: `attr` must implement IREE::Codegen::LayoutAttrInterface.
diff --git a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp
index 8552645..a467010 100644
--- a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp
+++ b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp
@@ -156,6 +156,32 @@
OpBuilder &builder,
RankedTensorType type,
ValueRange dynamicDims) const {
+ if (auto layoutsAttr = getLayouts()) {
+ if (llvm::any_of(layoutsAttr.getValue(), [](Attribute attr) {
+ return !llvm::isa<IREE::Encoding::EncodingLayoutAttrInterface>(attr);
+ })) {
+ return Value();
+ }
+
+ auto layoutsAttrArray =
+ llvm::to_vector_of<IREE::Encoding::EncodingLayoutAttrInterface>(
+ layoutsAttr.getValue());
+ Value res;
+ for (auto attr : layoutsAttrArray) {
+ Value requestedSize =
+ attr.calculateStorageSizeInBytes(loc, builder, type, dynamicDims);
+ if (!res) {
+ res = requestedSize;
+ continue;
+ }
+ res = builder.create<arith::MaxUIOp>(loc, res, requestedSize);
+ }
+ return res;
+ }
+
+ // TODO(hanchung): Deprecate the below logic once EncodingSpecialization pass
+ // is enabled by default. The layouts should be resolved and `roundDimsTo`
+ // will be deprecated.
SmallVector<int64_t> paddedShape(type.getShape());
SmallVector<Value> paddedDynamicDims(dynamicDims.begin(), dynamicDims.end());
ArrayRef<int64_t> roundDimsTo = getRoundDimsToArray();
diff --git a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingInterfaces.td b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingInterfaces.td
index 6eb994f..08fe49f 100644
--- a/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingInterfaces.td
+++ b/compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingInterfaces.td
@@ -28,7 +28,8 @@
InterfaceMethod<
/*desc=*/[{
Returns the storage size (in bytes) for the tensor types with an
- optional encoding.
+ optional encoding. Returns Value() if the size is unknown, i.e., it can
+ not be inferred with existing information.
}],
/*retTy=*/"::mlir::Value",
/*methodName=*/"calculateStorageSizeInBytes",
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors.mlir b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors.mlir
index 42764fe..ffe248a 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors.mlir
@@ -21,6 +21,25 @@
// -----
+#encoding_layout = #iree_cpu.vmvx_encoding_layout<configuration = {encoding_info = {innerDimsPos = [0, 1], innerTileSizes = [4, 16], outerDimsPerm = [0, 1]}}>
+#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], layouts = [#encoding_layout]>
+util.func public @sizeof_lhs_encoding_dynamic_using_layouts(%arg0: index, %arg1: index) -> index {
+ %0 = stream.tensor.sizeof tensor<?x?xf32, #encoding>{%arg0, %arg1} : index
+ util.return %0 : index
+}
+// CHECK-LABEL: @sizeof_lhs_encoding_dynamic_using_layouts
+// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK: %[[CEIL_DIV_D0:.+]] = arith.ceildivsi %arg0, %[[C4]]
+// CHECK: %[[PAD_D0:.+]] = arith.muli %[[CEIL_DIV_D0]], %[[C4]]
+// CHECK: %[[CEIL_DIV_D1:.+]] = arith.ceildivsi %arg1, %[[C16]]
+// CHECK: %[[PAD_D1:.+]] = arith.muli %[[CEIL_DIV_D1]], %[[C16]]
+// CHECK: %[[T0:.+]] = arith.muli %[[PAD_D0]], %[[C4]]
+// CHECK: %[[T1:.+]] = arith.muli %[[T0]], %[[PAD_D1]]
+// CHECK: return %[[T1]]
+
+// -----
+
#map = affine_map<(d0, d1, d2) -> (d0, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
@@ -42,6 +61,22 @@
// -----
+#encoding_layout = #iree_cpu.vmvx_encoding_layout<configuration = {encoding_info = {innerDimsPos = [0, 1], innerTileSizes = [4, 16], outerDimsPerm = [0, 1]}}>
+#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], layouts = [#encoding_layout]>
+util.func public @sizeof_lhs_encoding_partially_dynamic_using_layouts(%arg0: index) -> index {
+ %0 = stream.tensor.sizeof tensor<10x?xf32, #encoding>{%arg0} : index
+ util.return %0 : index
+}
+// CHECK-LABEL: @sizeof_lhs_encoding_partially_dynamic_using_layouts
+// CHECK-DAG: %[[C48:.+]] = arith.constant 48 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK: %[[CEIL_DIV_D1:.+]] = arith.ceildivsi %arg0, %[[C16]]
+// CHECK: %[[PAD_D1:.+]] = arith.muli %[[CEIL_DIV_D1]], %[[C16]]
+// CHECK: %[[T0:.+]] = arith.muli %[[PAD_D1]], %[[C48]]
+// CHECK: return %[[T0]]
+
+// -----
+
#map = affine_map<(d0, d1, d2) -> (d0, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
@@ -60,6 +95,29 @@
// -----
+// In GEMM, the RHS has the `(M, N, K) -> (K, N)` layout. The tile sizes
+// (i.e., [8, 16]) are for [dim_1, dim_0] in the encoding_info, where dim_1 is
+// N-dimension and dim_0 is K-dimension.
+#encoding_layout = #iree_cpu.vmvx_encoding_layout<configuration = {encoding_info = {innerDimsPos = [1, 0], innerTileSizes = [8, 16], outerDimsPerm = [1, 0]}}>
+#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], layouts = [#encoding_layout]>
+util.func public @sizeof_rhs_encoding_dynamic_using_layouts(%arg0: index, %arg1: index) -> index {
+ %0 = stream.tensor.sizeof tensor<?x?xf32, #encoding>{%arg0, %arg1} : index
+ util.return %0 : index
+}
+// CHECK-LABEL: @sizeof_rhs_encoding_dynamic_using_layouts
+// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index
+// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK: %[[CEIL_DIV_D1:.+]] = arith.ceildivsi %arg1, %[[C8]]
+// CHECK: %[[PAD_D1:.+]] = arith.muli %[[CEIL_DIV_D1]], %[[C8]]
+// CHECK: %[[CEIL_DIV_D0:.+]] = arith.ceildivsi %arg0, %[[C16]]
+// CHECK: %[[PAD_D0:.+]] = arith.muli %[[CEIL_DIV_D0]], %[[C16]]
+// CHECK: %[[T0:.+]] = arith.muli %[[PAD_D0]], %[[C4]]
+// CHECK: %[[T1:.+]] = arith.muli %[[T0]], %[[PAD_D1]]
+// CHECK: return %[[T1]]
+
+// -----
+
#map = affine_map<(d0, d1, d2) -> (d0, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
@@ -82,6 +140,25 @@
// -----
+#encoding_layout = #iree_cpu.vmvx_encoding_layout<configuration = {encoding_info = {innerDimsPos = [0, 1], innerTileSizes = [4, 8], outerDimsPerm = [0, 1]}}>
+#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], layouts = [#encoding_layout]>
+util.func public @sizeof_result_encoding_dynamic_using_layouts(%arg0: index, %arg1: index) -> index {
+ %0 = stream.tensor.sizeof tensor<?x?xf32, #encoding>{%arg0, %arg1} : index
+ util.return %0 : index
+}
+// CHECK-LABEL: @sizeof_result_encoding_dynamic_using_layouts
+// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index
+// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
+// CHECK: %[[CEIL_DIV_D0:.+]] = arith.ceildivsi %arg0, %[[C4]]
+// CHECK: %[[PAD_D0:.+]] = arith.muli %[[CEIL_DIV_D0]], %[[C4]]
+// CHECK: %[[CEIL_DIV_D1:.+]] = arith.ceildivsi %arg1, %[[C8]]
+// CHECK: %[[PAD_D1:.+]] = arith.muli %[[CEIL_DIV_D1]], %[[C8]]
+// CHECK: %[[T0:.+]] = arith.muli %[[PAD_D0]], %[[C4]]
+// CHECK: %[[T1:.+]] = arith.muli %[[T0]], %[[PAD_D1]]
+// CHECK: return %[[T1]]
+
+// -----
+
#map = affine_map<(d0, d1, d2) -> (d0, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
@@ -103,6 +180,30 @@
// -----
+// The layout is as the same as the the matmul LHS layout because it broadcasts
+// across the batch dimension. The test is preserved for having the same test
+// suite of non-layouts style encoding. I.e., this is the resolved layout
+// version of the below sizeof_lhs_encoding_with_bcast_across_batch_dim_dynamic
+// test.
+#encoding_layout = #iree_cpu.vmvx_encoding_layout<configuration = {encoding_info = {innerDimsPos = [0, 1], innerTileSizes = [4, 16], outerDimsPerm = [0, 1]}}>
+#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], layouts = [#encoding_layout]>
+util.func public @sizeof_lhs_encoding_with_bcast_across_batch_dim_dynamic_using_layouts(%arg0: index, %arg1: index) -> index {
+ %0 = stream.tensor.sizeof tensor<?x?xf32, #encoding>{%arg0, %arg1} : index
+ util.return %0 : index
+}
+// CHECK-LABEL: @sizeof_lhs_encoding_with_bcast_across_batch_dim_dynamic_using_layouts
+// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK: %[[CEIL_DIV_D0:.+]] = arith.ceildivsi %arg0, %[[C4]]
+// CHECK: %[[PAD_D0:.+]] = arith.muli %[[CEIL_DIV_D0]], %[[C4]]
+// CHECK: %[[CEIL_DIV_D1:.+]] = arith.ceildivsi %arg1, %[[C16]]
+// CHECK: %[[PAD_D1:.+]] = arith.muli %[[CEIL_DIV_D1]], %[[C16]]
+// CHECK: %[[T0:.+]] = arith.muli %[[PAD_D0]], %[[C4]]
+// CHECK: %[[T1:.+]] = arith.muli %[[T0]], %[[PAD_D1]]
+// CHECK: return %[[T1]]
+
+// -----
+
#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
@@ -125,6 +226,28 @@
// -----
+// The M-dimension inner tile is not present because it broadcasts across the
+// M-dimension. We do not need to pack the M-dimension in this case.
+#encoding_layout = #iree_cpu.vmvx_encoding_layout<configuration = {encoding_info = {innerDimsPos = [1], innerTileSizes = [16], outerDimsPerm = [0, 1]}}>
+#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], layouts = [#encoding_layout]>
+util.func public @sizeof_lhs_encoding_with_bcast_across_m_dim_dynamic_using_layouts(%arg0: index, %arg1: index) -> index {
+ %0 = stream.tensor.sizeof tensor<?x?xf32, #encoding>{%arg0, %arg1} : index
+ util.return %0 : index
+}
+// CHECK-LABEL: @sizeof_lhs_encoding_with_bcast_across_m_dim_dynamic_using_layouts
+// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+// CHECK: %[[CEIL_DIV_D1:.+]] = arith.ceildivsi %arg1, %[[C16]]
+// CHECK: %[[PAD_D1:.+]] = arith.muli %[[CEIL_DIV_D1]], %[[C16]]
+//
+// Multiplied by 4 because f32 has 4 bytes.
+//
+// CHECK: %[[T0:.+]] = arith.muli %arg0, %[[C4]]
+// CHECK: %[[T1:.+]] = arith.muli %[[T0]], %[[PAD_D1]]
+// CHECK: return %[[T1]]
+
+// -----
+
#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
@@ -148,6 +271,44 @@
// -----
+#encoding_layout_0 = #iree_cpu.cpu_encoding_layout<configuration = {encoding_info = {innerDimsPos = [0, 1], innerTileSizes = [4, 8], outerDimsPerm = [0, 1]}}>
+#encoding_layout_1 = #iree_cpu.vmvx_encoding_layout<configuration = {encoding_info = {innerDimsPos = [0, 1], innerTileSizes = [2, 16], outerDimsPerm = [0, 1]}}>
+#encoding = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], layouts = [#encoding_layout_0, #encoding_layout_1]>
+util.func public @sizeof_multi_encoding_layouts(%arg0: index, %arg1: index) -> index {
+ %0 = stream.tensor.sizeof tensor<?x?xf32, #encoding>{%arg0, %arg1} : index
+ util.return %0 : index
+}
+// CHECK-LABEL: @sizeof_multi_encoding_layouts
+// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
+// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index
+// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
+// CHECK-DAG: %[[C16:.+]] = arith.constant 16 : index
+//
+// Check for the first layout.
+//
+// CHECK: %[[CEIL_DIV_D0:.+]] = arith.ceildivsi %arg0, %[[C4]]
+// CHECK: %[[PAD_D0:.+]] = arith.muli %[[CEIL_DIV_D0]], %[[C4]]
+// CHECK: %[[CEIL_DIV_D1:.+]] = arith.ceildivsi %arg1, %[[C8]]
+// CHECK: %[[PAD_D1:.+]] = arith.muli %[[CEIL_DIV_D1]], %[[C8]]
+// CHECK: %[[T0:.+]] = arith.muli %[[PAD_D0]], %[[C4]]
+// CHECK: %[[SIZE0:.+]] = arith.muli %[[T0]], %[[PAD_D1]]
+//
+// Check for the first layout.
+//
+// CHECK: %[[CEIL_DIV_D0:.+]] = arith.ceildivsi %arg0, %[[C2]]
+// CHECK: %[[PAD_D0:.+]] = arith.muli %[[CEIL_DIV_D0]], %[[C2]]
+// CHECK: %[[CEIL_DIV_D1:.+]] = arith.ceildivsi %arg1, %[[C16]]
+// CHECK: %[[PAD_D1:.+]] = arith.muli %[[CEIL_DIV_D1]], %[[C16]]
+// CHECK: %[[T1:.+]] = arith.muli %[[PAD_D0]], %[[C4]]
+// CHECK: %[[SIZE1:.+]] = arith.muli %[[T1]], %[[PAD_D1]]
+//
+// Return the max value.
+//
+// CHECK: %[[RES:.+]] = arith.maxui %[[SIZE0]], %[[SIZE1]]
+// CHECK: return %[[RES]]
+
+// -----
+
// CHECK-LABEL: @denseTensorEmpty
util.func public @denseTensorEmpty(%arg0: index, %arg1: index) -> !stream.resource<*> {
// CHECK: %[[RET:.+]] = stream.async.alloca : !stream.resource<*>{%arg1}
