Introduce `iree_codegen.query_tile_sizes` replacing `vmvx.query_tile_sizes` (#14147)
It's still only used by VMVX (as that's the only place we can't resolve
tile sizes at compile time) but it's a step towards removing legacy
VMVXOps.
diff --git a/compiler/src/iree/compiler/Codegen/Common/EncodingInfo.h b/compiler/src/iree/compiler/Codegen/Common/EncodingInfo.h
index 9691993..f5bf182 100644
--- a/compiler/src/iree/compiler/Codegen/Common/EncodingInfo.h
+++ b/compiler/src/iree/compiler/Codegen/Common/EncodingInfo.h
@@ -8,34 +8,18 @@
#define IREE_COMPILER_SRC_IREE_COMPILER_CODEGEN_COMMON_ENCODINGINFO_H_
#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h"
-#include "iree/compiler/Codegen/Utils/EncodingInfo.h"
+#include "iree/compiler/Codegen/Utils/EncodingUtils.h"
#include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
namespace mlir {
namespace iree_compiler {
-enum class MatmulOperandRole {
- LHS,
- RHS,
- RESULT,
-};
-
struct MatmulTileParams {
int64_t M = 1;
int64_t K = 1;
int64_t N = 1;
};
-/// Extracts encoding from the `tensorType` if specified.
-std::optional<IREE::LinalgExt::TensorEncoding> getEncoding(
- RankedTensorType tensorType);
-
-std::optional<MatmulType> getMatmulType(
- IREE::LinalgExt::TensorEncoding encoding);
-
-std::optional<MatmulOperandRole> getMatmulOperandRole(
- IREE::LinalgExt::TensorEncoding encoding);
-
void adjustTileSizesToNarrowStaticShape(
IREE::LinalgExt::MaterializeEncodingInfo &encodingInfo,
ArrayRef<int64_t> shape);
diff --git a/compiler/src/iree/compiler/Codegen/Common/LowerUKernelsToCalls.cpp b/compiler/src/iree/compiler/Codegen/Common/LowerUKernelsToCalls.cpp
index 175b64d..023b889 100644
--- a/compiler/src/iree/compiler/Codegen/Common/LowerUKernelsToCalls.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/LowerUKernelsToCalls.cpp
@@ -7,6 +7,7 @@
#include "iree/compiler/Codegen/Common/CommonPasses.h"
#include "iree/compiler/Codegen/Interfaces/UKernelOpInterface.h"
#include "iree/compiler/Codegen/PassDetail.h"
+#include "llvm/ADT/MapVector.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -27,7 +28,7 @@
void LowerUKernelOpsToCallsPass::runOnOperation() {
MLIRContext *context = &getContext();
RewritePatternSet patterns(context);
- SmallVector<Operation *> toDelete;
+ llvm::MapVector<IREE::Codegen::UKernelOpInterface, func::CallOp> toReplace;
Operation *errorOp = nullptr;
IRRewriter rewriter(context);
WalkResult result = getOperation().walk(
@@ -40,7 +41,7 @@
errorOp = microKernelOp;
return WalkResult::interrupt();
}
- toDelete.push_back(microKernelOp);
+ toReplace[microKernelOp] = callOp.value();
return WalkResult::advance();
});
if (result.wasInterrupted()) {
@@ -48,8 +49,8 @@
"failed to lower micro kernel operation to function call");
return signalPassFailure();
}
- for (auto op : toDelete) {
- op->erase();
+ for (auto r : toReplace) {
+ rewriter.replaceOp(r.first, r.second->getResults());
}
}
diff --git a/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoPackUnPack.cpp b/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoPackUnPack.cpp
index 22aa0b5..0958fb4 100644
--- a/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoPackUnPack.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoPackUnPack.cpp
@@ -12,6 +12,7 @@
#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h"
#include "iree/compiler/Codegen/Common/CommonPasses.h"
#include "iree/compiler/Codegen/Common/EncodingInfo.h"
+#include "iree/compiler/Codegen/Dialect/IREECodegenOps.h"
#include "iree/compiler/Codegen/PassDetail.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
@@ -270,44 +271,6 @@
return encodingInfo;
}
-std::optional<TensorEncoding> getEncoding(RankedTensorType tensorType) {
- auto encodingAttr =
- llvm::dyn_cast_if_present<EncodingAttr>(tensorType.getEncoding());
- if (!encodingAttr) return std::nullopt;
- return encodingAttr.getEncoding().getValue();
-}
-
-std::optional<MatmulType> getMatmulType(TensorEncoding encoding) {
- switch (encoding) {
- case TensorEncoding::MATMUL_F32F32F32_LHS:
- case TensorEncoding::MATMUL_F32F32F32_RHS:
- case TensorEncoding::MATMUL_F32F32F32_RESULT:
- return MatmulType::F32F32F32;
- case TensorEncoding::MATMUL_I8I8I32_LHS:
- case TensorEncoding::MATMUL_I8I8I32_RHS:
- case TensorEncoding::MATMUL_I8I8I32_RESULT:
- return MatmulType::I8I8I32;
- default:
- return std::nullopt;
- }
-}
-
-std::optional<MatmulOperandRole> getMatmulOperandRole(TensorEncoding encoding) {
- switch (encoding) {
- case TensorEncoding::MATMUL_F32F32F32_LHS:
- case TensorEncoding::MATMUL_I8I8I32_LHS:
- return MatmulOperandRole::LHS;
- case TensorEncoding::MATMUL_F32F32F32_RHS:
- case TensorEncoding::MATMUL_I8I8I32_RHS:
- return MatmulOperandRole::RHS;
- case TensorEncoding::MATMUL_F32F32F32_RESULT:
- case TensorEncoding::MATMUL_I8I8I32_RESULT:
- return MatmulOperandRole::RESULT;
- default:
- return std::nullopt;
- }
-}
-
void adjustTileSizesToNarrowStaticShape(MaterializeEncodingInfo &encodingInfo,
ArrayRef<int64_t> shape) {
for (size_t i = 0; i < shape.size(); i++) {
@@ -334,38 +297,11 @@
FailureOr<MaterializeEncodingValueInfo>
chooseDynamicEncodingInfoVMVXMicrokernels(RankedTensorType tensorType,
OpBuilder &builder, Location loc) {
- std::optional<TensorEncoding> encoding = getEncoding(tensorType);
- if (!encoding) return failure();
- auto matmulType = getMatmulType(*encoding);
- auto matmulOperandRole = getMatmulOperandRole(*encoding);
- if (!matmulType || !matmulOperandRole) return failure();
- uint32_t flags = 0;
- if (*matmulType == MatmulType::F32F32F32) {
- flags |= IREE_UK_FLAG_QUERY_TILE_SIZES_OPERATION_MATMUL_F32F32F32;
- } else if (*matmulType == MatmulType::I8I8I32) {
- flags |= IREE_UK_FLAG_QUERY_TILE_SIZES_OPERATION_MATMUL_I8I8I32;
- } else {
- return failure();
- }
- if (*matmulOperandRole == MatmulOperandRole::LHS) {
- flags |= IREE_UK_FLAG_QUERY_TILE_SIZES_OPERAND_ROLE_LHS;
- } else if (*matmulOperandRole == MatmulOperandRole::RHS) {
- flags |= IREE_UK_FLAG_QUERY_TILE_SIZES_OPERAND_ROLE_RHS;
- } else if (*matmulOperandRole == MatmulOperandRole::RESULT) {
- flags |= IREE_UK_FLAG_QUERY_TILE_SIZES_OPERAND_ROLE_RESULT;
- } else {
- return failure();
- }
- SmallVector<Type> tileSizesTypes(tensorType.getRank(),
- builder.getIndexType());
- SmallVector<Value> shapeValues;
- for (int64_t i : tensorType.getShape()) {
- shapeValues.push_back(builder.create<arith::ConstantIndexOp>(loc, i));
- }
- auto op = builder.create<IREE::VMVX::QueryTileSizesOp>(
- loc, tileSizesTypes, shapeValues, builder.getI32IntegerAttr(flags));
+ SmallVector<Type> resultTypes(tensorType.getRank(), builder.getIndexType());
+ auto op = builder.create<IREE::Codegen::QueryTileSizesOp>(
+ loc, resultTypes, TypeAttr::get(tensorType));
MaterializeEncodingValueInfo result;
- result.innerTileSizes = op.getTileSizes();
+ result.innerTileSizes = op.getResults();
return result;
}
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/convert_to_destination_passing_style.mlir b/compiler/src/iree/compiler/Codegen/Common/test/convert_to_destination_passing_style.mlir
index 524f13f..4454fd6 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/convert_to_destination_passing_style.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/convert_to_destination_passing_style.mlir
@@ -751,15 +751,11 @@
%c512 = arith.constant 512 : index
%c0 = arith.constant 0 : index
%c16 = arith.constant 16 : index
- %0:2 = vmvx.query_tile_sizes sizes(%c16, %c16) flags(1245184) -> index, index
+ %0:2 = iree_codegen.query_tile_sizes tensor<16x16xi32, #iree_linalg_ext.encoding<MATMUL_I8I8I32_RESULT>> -> index, index
%1 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%0#0]
%2 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%0#1]
%3 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c512) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi32>>{%1, %2, %0#0, %0#1}
%4 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<writeonly:tensor<1x1xi32>>
- %5:2 = vmvx.query_tile_sizes sizes(%c16, %c16) flags(1245184) -> index, index
- %6 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%5#0]
- %7 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%5#1]
- %8:2 = vmvx.query_tile_sizes sizes(%c16, %c16) flags(1245184) -> index, index
%workgroup_id_x = hal.interface.workgroup.id[0] : index
%workgroup_count_x = hal.interface.workgroup.count[0] : index
%workgroup_id_y = hal.interface.workgroup.id[1] : index
@@ -770,13 +766,13 @@
%11 = affine.apply affine_map<()[s0] -> (s0 * 16)>()[%workgroup_id_x]
%12 = affine.apply affine_map<()[s0] -> (s0 * 16)>()[%workgroup_count_x]
scf.for %arg1 = %11 to %c1 step %12 {
- %13 = affine.apply affine_map<(d0)[s0] -> (d0 mod s0)>(%arg0)[%8#0]
- %14 = affine.apply affine_map<(d0)[s0] -> (d0 mod s0)>(%arg1)[%8#1]
- %15 = affine.apply affine_map<(d0)[s0] -> (d0 floordiv s0)>(%arg0)[%8#0]
- %16 = affine.apply affine_map<(d0)[s0] -> (d0 floordiv s0)>(%arg1)[%8#1]
- %17 = flow.dispatch.tensor.load %3, offsets = [%15, %16, 0, 0], sizes = [%c1, %c1, %8#0, %8#1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi32>>{%6, %7, %5#0, %5#1} -> tensor<?x?x?x?xi32>
- %18 = tensor.empty(%8#0, %8#1) : tensor<?x?xi32>
- %19 = tensor.unpack %17 inner_dims_pos = [0, 1] inner_tiles = [%8#0, %8#1] into %18 : tensor<?x?x?x?xi32> -> tensor<?x?xi32>
+ %13 = affine.apply affine_map<(d0)[s0] -> (d0 mod s0)>(%arg0)[%0#0]
+ %14 = affine.apply affine_map<(d0)[s0] -> (d0 mod s0)>(%arg1)[%0#1]
+ %15 = affine.apply affine_map<(d0)[s0] -> (d0 floordiv s0)>(%arg0)[%0#0]
+ %16 = affine.apply affine_map<(d0)[s0] -> (d0 floordiv s0)>(%arg1)[%0#1]
+ %17 = flow.dispatch.tensor.load %3, offsets = [%15, %16, 0, 0], sizes = [%c1, %c1, %0#0, %0#1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi32>>{%1, %2, %0#0, %0#1} -> tensor<?x?x?x?xi32>
+ %18 = tensor.empty(%0#0, %0#1) : tensor<?x?xi32>
+ %19 = tensor.unpack %17 inner_dims_pos = [0, 1] inner_tiles = [%0#0, %0#1] into %18 : tensor<?x?x?x?xi32> -> tensor<?x?xi32>
%extracted_slice = tensor.extract_slice %19[%13, %14] [1, 1] [1, 1] : tensor<?x?xi32> to tensor<1x1xi32>
%cast = tensor.cast %extracted_slice : tensor<1x1xi32> to tensor<?x?xi32>
flow.dispatch.tensor.store %cast, %4, offsets = [%arg0, %arg1], sizes = [%c1, %c1], strides = [1, 1] : tensor<?x?xi32> -> !flow.dispatch.tensor<writeonly:tensor<1x1xi32>>
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/tile_and_distribute_to_workgroups.mlir b/compiler/src/iree/compiler/Codegen/Common/test/tile_and_distribute_to_workgroups.mlir
index d1050a6..c19fed9 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/tile_and_distribute_to_workgroups.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/tile_and_distribute_to_workgroups.mlir
@@ -2242,23 +2242,19 @@
%cst = arith.constant dense<[-918, -4433, 87, -234, -21393, 7738, 529, -8835, -16817, -375, -199, 572, 5082, 15569, -186, 4955]> : tensor<16xi32>
%c12544 = arith.constant 12544 : index
%c16 = arith.constant 16 : index
- %0:2 = vmvx.query_tile_sizes sizes(%c12544, %c16) flags(1245184) -> index, index
+ %0:2 = iree_codegen.query_tile_sizes tensor<12544x16xi32, #iree_linalg_ext.encoding<MATMUL_I8I8I32_RESULT>> -> index, index
%1 = affine.apply affine_map<()[s0] -> (12544 ceildiv s0)>()[%0#0]
%2 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%0#1]
%3 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c200960) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi32>>{%1, %2, %0#0, %0#1}
%4 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c1003776) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<12544xi32>>
%5 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c1053952) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<16xi32>>
%6 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<writeonly:tensor<12544x16xi32>>
- %7:2 = vmvx.query_tile_sizes sizes(%c12544, %c16) flags(1245184) -> index, index
- %8 = affine.apply affine_map<()[s0] -> (12544 ceildiv s0)>()[%7#0]
- %9 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%7#1]
- %10 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0], sizes = [%8, %9, %7#0, %7#1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi32>>{%8, %9, %7#0, %7#1} -> tensor<?x?x?x?xi32>
+ %10 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0], sizes = [%1, %2, %0#0, %0#1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi32>>{%1, %2, %0#0, %0#1} -> tensor<?x?x?x?xi32>
%11 = flow.dispatch.tensor.load %4, offsets = [0], sizes = [12544], strides = [1] : !flow.dispatch.tensor<readonly:tensor<12544xi32>> -> tensor<12544xi32>
%12 = flow.dispatch.tensor.load %5, offsets = [0], sizes = [16], strides = [1] : !flow.dispatch.tensor<readonly:tensor<16xi32>> -> tensor<16xi32>
%13 = tensor.empty() : tensor<12544x16xi32>
%14 = tensor.empty() : tensor<12544x16xi32>
- %15:2 = vmvx.query_tile_sizes sizes(%c12544, %c16) flags(1245184) -> index, index
- %16 = tensor.unpack %10 inner_dims_pos = [0, 1] inner_tiles = [%15#0, %15#1] into %14 {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[16, 16]]>} : tensor<?x?x?x?xi32> -> tensor<12544x16xi32>
+ %16 = tensor.unpack %10 inner_dims_pos = [0, 1] inner_tiles = [%0#0, %0#1] into %14 {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[16, 16]]>} : tensor<?x?x?x?xi32> -> tensor<12544x16xi32>
%17 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%cst, %16, %11, %12 : tensor<16xi32>, tensor<12544x16xi32>, tensor<12544xi32>, tensor<16xi32>) outs(%13 : tensor<12544x16xi32>) {
^bb0(%in: i32, %in_0: i32, %in_1: i32, %in_2: i32, %out: i32):
%18 = arith.muli %in_1, %c-30_i32 : i32
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Dialect/BUILD.bazel
index 4383958..49181a9 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Dialect/BUILD.bazel
@@ -16,6 +16,7 @@
exports_files([
"IREECodegenAttributes.td",
"IREECodegenDialect.td",
+ "IREECodegenOps.td",
"LoweringConfig.td",
"UKernelOps.td",
])
@@ -26,6 +27,7 @@
[
"IREECodegenAttributes.td",
"IREECodegenDialect.td",
+ "IREECodegenOps.td",
"LoweringConfig.td",
"UKernelOps.td",
],
@@ -35,6 +37,7 @@
"//compiler/src/iree/compiler/Codegen/Interfaces:td_files",
"@llvm-project//mlir:DestinationStyleOpInterfaceTdFiles",
"@llvm-project//mlir:OpBaseTdFiles",
+ "@llvm-project//mlir:SideEffectInterfacesTdFiles",
],
)
@@ -42,17 +45,21 @@
name = "IREECodegenDialect",
srcs = [
"IREECodegenDialect.cpp",
+ "IREECodegenOps.cpp",
"LoweringConfig.cpp",
"UKernelOps.cpp",
],
hdrs = [
"IREECodegenDialect.h",
+ "IREECodegenOps.h",
"LoweringConfig.h",
"UKernelOps.h",
],
textual_hdrs = [
"IREECodegenDialect.cpp.inc",
"IREECodegenDialect.h.inc",
+ "IREECodegenOps.cpp.inc",
+ "IREECodegenOps.h.inc",
"LoweringConfig.cpp.inc",
"LoweringConfig.h.inc",
"LoweringConfigEnums.cpp.inc",
@@ -62,6 +69,7 @@
],
deps = [
":IREECodegenDialectGen",
+ ":IREECodegenOpsGen",
":LoweringConfigGen",
":UKernelOpsGen",
"//compiler/src/iree/compiler/Codegen/Interfaces:UKernelOpInterface",
@@ -99,6 +107,26 @@
)
iree_gentbl_cc_library(
+ name = "IREECodegenOpsGen",
+ tbl_outs = [
+ (
+ ["--gen-op-decls"],
+ "IREECodegenOps.h.inc",
+ ),
+ (
+ ["--gen-op-defs"],
+ "IREECodegenOps.cpp.inc",
+ ),
+ ],
+ tblgen = "@llvm-project//mlir:mlir-tblgen",
+ td_file = "IREECodegenOps.td",
+ deps = [
+ ":td_files",
+ "//compiler/src/iree/compiler/Codegen/Interfaces:td_files",
+ ],
+)
+
+iree_gentbl_cc_library(
name = "LoweringConfigGen",
tbl_outs = [
(
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Dialect/CMakeLists.txt
index ce421c2..7efe4f0 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Dialect/CMakeLists.txt
@@ -15,11 +15,14 @@
IREECodegenDialect
HDRS
"IREECodegenDialect.h"
+ "IREECodegenOps.h"
"LoweringConfig.h"
"UKernelOps.h"
TEXTUAL_HDRS
"IREECodegenDialect.cpp.inc"
"IREECodegenDialect.h.inc"
+ "IREECodegenOps.cpp.inc"
+ "IREECodegenOps.h.inc"
"LoweringConfig.cpp.inc"
"LoweringConfig.h.inc"
"LoweringConfigEnums.cpp.inc"
@@ -28,10 +31,12 @@
"UKernelOps.h.inc"
SRCS
"IREECodegenDialect.cpp"
+ "IREECodegenOps.cpp"
"LoweringConfig.cpp"
"UKernelOps.cpp"
DEPS
::IREECodegenDialectGen
+ ::IREECodegenOpsGen
::LoweringConfigGen
::UKernelOpsGen
LLVMSupport
@@ -62,6 +67,16 @@
iree_tablegen_library(
NAME
+ IREECodegenOpsGen
+ TD_FILE
+ "IREECodegenOps.td"
+ OUTS
+ --gen-op-decls IREECodegenOps.h.inc
+ --gen-op-defs IREECodegenOps.cpp.inc
+)
+
+iree_tablegen_library(
+ NAME
LoweringConfigGen
TD_FILE
"LoweringConfig.td"
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/IREECodegenDialect.cpp b/compiler/src/iree/compiler/Codegen/Dialect/IREECodegenDialect.cpp
index 11178b4..6f6b58a 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/IREECodegenDialect.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/IREECodegenDialect.cpp
@@ -7,6 +7,7 @@
#include "iree/compiler/Codegen/Dialect/IREECodegenDialect.h"
#include "iree/compiler/Codegen/Dialect/IREECodegenDialect.cpp.inc"
+#include "iree/compiler/Codegen/Dialect/IREECodegenOps.h"
#include "iree/compiler/Codegen/Dialect/LoweringConfig.h"
#include "iree/compiler/Codegen/Dialect/UKernelOps.h"
#include "mlir/IR/DialectImplementation.h"
@@ -39,6 +40,10 @@
addOperations<
#define GET_OP_LIST
+#include "iree/compiler/Codegen/Dialect/IREECodegenOps.cpp.inc"
+ >();
+ addOperations<
+#define GET_OP_LIST
#include "iree/compiler/Codegen/Dialect/UKernelOps.cpp.inc"
>();
}
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/IREECodegenOps.cpp b/compiler/src/iree/compiler/Codegen/Dialect/IREECodegenOps.cpp
new file mode 100644
index 0000000..f917e9a
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Dialect/IREECodegenOps.cpp
@@ -0,0 +1,32 @@
+// Copyright 2023 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/IREECodegenOps.h"
+
+#include "iree/compiler/Codegen/Dialect/IREECodegenDialect.h"
+#include "iree/compiler/Codegen/Utils/EncodingUtils.h"
+#include "iree/compiler/Codegen/Utils/Utils.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Bufferization/IR/DstBufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/Support/LLVM.h"
+
+// clang-format off
+#define GET_OP_CLASSES
+#include "iree/compiler/Codegen/Dialect/IREECodegenOps.cpp.inc" // IWYU pragma: keep
+// clang-format on
+
+namespace mlir {
+namespace iree_compiler {
+namespace IREE {
+namespace Codegen {} // namespace Codegen
+} // namespace IREE
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/IREECodegenOps.h b/compiler/src/iree/compiler/Codegen/Dialect/IREECodegenOps.h
new file mode 100644
index 0000000..926e576
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Dialect/IREECodegenOps.h
@@ -0,0 +1,21 @@
+// Copyright 2023 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_IREECODEGENOPS_H_
+#define IREE_COMPILER_CODEGEN_DIALECT_IREECODEGENOPS_H_
+
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+// clang-format off
+#define GET_OP_CLASSES
+#include "iree/compiler/Codegen/Dialect/IREECodegenOps.h.inc" // IWYU pragma: export
+// clang-format on
+
+#endif // #ifndef IREE_COMPILER_CODEGEN_DIALECT_IREECODEGENOPS_H_
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/IREECodegenOps.td b/compiler/src/iree/compiler/Codegen/Dialect/IREECodegenOps.td
new file mode 100644
index 0000000..6202ced
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Dialect/IREECodegenOps.td
@@ -0,0 +1,31 @@
+// Copyright 2023 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_IREECODEGENOPS
+#define IREE_CODEGEN_DIALECT_IREECODEGENOPS
+
+include "iree/compiler/Codegen/Dialect/IREECodegenDialect.td"
+include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+def TensorTypeAttr : TypeAttrBase<"TensorType", "Tensor type attribute">;
+
+def IREECodegen_QueryTileSizesOp :
+ Op<IREECodegen_Dialect, "query_tile_sizes", [Pure]> {
+ let summary = "Query tile sizes";
+
+ let description = [{
+ Query tile sizes
+ }];
+
+ let arguments = (ins TensorTypeAttr:$tensor_type);
+ let results = (outs Variadic<Index>:$results);
+ let assemblyFormat = [{
+ attr-dict $tensor_type `->` type($results)
+ }];
+}
+
+#endif // IREE_CODEGEN_DIALECT_IREECODEGENOPS
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/UKernelOps.cpp b/compiler/src/iree/compiler/Codegen/Dialect/UKernelOps.cpp
index 4f741b6..4addd65 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/UKernelOps.cpp
+++ b/compiler/src/iree/compiler/Codegen/Dialect/UKernelOps.cpp
@@ -8,7 +8,7 @@
#include "iree/builtins/ukernel/exported_bits.h"
#include "iree/compiler/Codegen/Dialect/IREECodegenDialect.h"
-#include "iree/compiler/Codegen/Utils/EncodingInfo.h"
+#include "iree/compiler/Codegen/Utils/EncodingUtils.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/FormatVariadic.h"
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/UKernelOps.td b/compiler/src/iree/compiler/Codegen/Dialect/UKernelOps.td
index ca35b92..3cbacc6 100644
--- a/compiler/src/iree/compiler/Codegen/Dialect/UKernelOps.td
+++ b/compiler/src/iree/compiler/Codegen/Dialect/UKernelOps.td
@@ -64,7 +64,7 @@
Variadic<AnyType>:$other_operands,
OptionalAttr<DictionaryAttr>:$fn_def_attrs,
OptionalAttr<IndexAttr>:$strided_outer_dims);
- let results = (outs Variadic<AnyRankedTensor>:$results);
+ let results = (outs Variadic<AnyType>:$results);
let assemblyFormat = [{
attr-dict $u_kernel_fn_name
(`ins` `(` $inputs^ `:` type($inputs) `)`)?
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerToUKernels.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerToUKernels.cpp
index 8c0442e..c34216c 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerToUKernels.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPULowerToUKernels.cpp
@@ -6,15 +6,21 @@
#include "iree/builtins/ukernel/exported_bits.h"
#include "iree/compiler/Codegen/Dialect/IREECodegenDialect.h"
+#include "iree/compiler/Codegen/Dialect/IREECodegenOps.h"
#include "iree/compiler/Codegen/Dialect/UKernelOps.h"
#include "iree/compiler/Codegen/LLVMCPU/LLVMCPUPasses.h"
#include "iree/compiler/Codegen/PassDetail.h"
-#include "iree/compiler/Codegen/Utils/EncodingInfo.h"
+#include "iree/compiler/Codegen/Utils/EncodingUtils.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Matchers.h"
+#include "mlir/IR/TypeRange.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-
namespace mlir {
namespace iree_compiler {
@@ -327,6 +333,60 @@
genericMicroKernelOp.getOperation());
}
+static FailureOr<IREE::Codegen::UKernelOpInterface> matchDAGForUKernel(
+ RewriterBase &rewriter, IREE::Codegen::QueryTileSizesOp op) {
+ auto tensorType = op.getTensorType().dyn_cast<RankedTensorType>();
+ if (!tensorType) {
+ return rewriter.notifyMatchFailure(op,
+ "need a ranked tensor type attribute");
+ }
+ if (tensorType.getRank() != 2) {
+ return rewriter.notifyMatchFailure(op, "only the 2D case is implemented");
+ }
+ auto encoding = getEncoding(tensorType);
+ if (!encoding) {
+ return rewriter.notifyMatchFailure(op, "no TensorEncoding attribute");
+ }
+ auto matmulType = getMatmulType(*encoding);
+ auto matmulOperandRole = getMatmulOperandRole(*encoding);
+ if (!matmulType || !matmulOperandRole) {
+ return rewriter.notifyMatchFailure(op, "unhandled TensorEncoding");
+ }
+ uint32_t flags = 0;
+ if (*matmulType == MatmulType::F32F32F32) {
+ flags |= IREE_UK_FLAG_QUERY_TILE_SIZES_OPERATION_MATMUL_F32F32F32;
+ } else if (*matmulType == MatmulType::I8I8I32) {
+ flags |= IREE_UK_FLAG_QUERY_TILE_SIZES_OPERATION_MATMUL_I8I8I32;
+ } else {
+ return failure();
+ }
+ if (*matmulOperandRole == MatmulOperandRole::LHS) {
+ flags |= IREE_UK_FLAG_QUERY_TILE_SIZES_OPERAND_ROLE_LHS;
+ } else if (*matmulOperandRole == MatmulOperandRole::RHS) {
+ flags |= IREE_UK_FLAG_QUERY_TILE_SIZES_OPERAND_ROLE_RHS;
+ } else if (*matmulOperandRole == MatmulOperandRole::RESULT) {
+ flags |= IREE_UK_FLAG_QUERY_TILE_SIZES_OPERAND_ROLE_RESULT;
+ } else {
+ return failure();
+ }
+ SmallVector<Type> resultTypes(tensorType.getRank(), rewriter.getIndexType());
+ SmallVector<Value> inputValues;
+ Location loc = op.getLoc();
+ for (int64_t i : tensorType.getShape()) {
+ inputValues.push_back(rewriter.create<arith::ConstantIndexOp>(loc, i));
+ }
+ inputValues.push_back(rewriter.create<arith::ConstantIntOp>(loc, flags, 32));
+ auto targetAttr = IREE::HAL::ExecutableTargetAttr::lookup(op);
+ auto fn = getFnNameAndDefAttrs("query_tile_sizes.2d", rewriter, targetAttr);
+ auto genericMicroKernelOp = rewriter.create<IREE::Codegen::UKernelGenericOp>(
+ loc, resultTypes, fn.name, inputValues, /*outs=*/ValueRange{},
+ /*other_operands=*/ValueRange{},
+ /*fn_def_attrs=*/rewriter.getDictionaryAttr(fn.defAttrs),
+ /*strided_outer_dims=*/IntegerAttr{});
+ return cast<IREE::Codegen::UKernelOpInterface>(
+ genericMicroKernelOp.getOperation());
+}
+
namespace {
template <typename OpType>
@@ -335,10 +395,6 @@
LogicalResult matchAndRewrite(OpType op,
PatternRewriter &rewriter) const override {
- if (!op.hasTensorSemantics()) {
- return rewriter.notifyMatchFailure(
- op, "operation needs to have tensor semantics");
- }
FailureOr<IREE::Codegen::UKernelOpInterface> ukernelOp =
matchDAGForUKernel(rewriter, op);
if (failed(ukernelOp)) {
@@ -357,7 +413,9 @@
RewritePatternSet patterns(context);
patterns.insert<LowerToUKernelPattern<linalg::Mmt4DOp>,
LowerToUKernelPattern<tensor::PackOp>,
- LowerToUKernelPattern<tensor::UnPackOp>>(context);
+ LowerToUKernelPattern<tensor::UnPackOp>,
+ LowerToUKernelPattern<IREE::Codegen::QueryTileSizesOp>>(
+ context);
if (failed(
applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)))) {
return signalPassFailure();
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUMaterializeEncodingPass.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUMaterializeEncodingPass.cpp
index 51fadcc..2e3a0fe 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUMaterializeEncodingPass.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUMaterializeEncodingPass.cpp
@@ -7,6 +7,7 @@
#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtDialect.h"
#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h"
#include "iree/compiler/Codegen/Common/EncodingInfo.h"
+#include "iree/compiler/Codegen/Dialect/IREECodegenDialect.h"
#include "iree/compiler/Codegen/LLVMCPU/Utils.h"
#include "iree/compiler/Codegen/PassDetail.h"
#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
@@ -100,9 +101,10 @@
struct LLVMCPUMaterializeEncodingPass
: public LLVMCPUMaterializeEncodingBase<LLVMCPUMaterializeEncodingPass> {
void getDependentDialects(DialectRegistry ®istry) const override {
- registry.insert<arith::ArithDialect, affine::AffineDialect,
- IREE::Flow::FlowDialect,
- IREE::LinalgExt::IREELinalgExtDialect>();
+ registry
+ .insert<arith::ArithDialect, affine::AffineDialect,
+ IREE::Flow::FlowDialect, IREE::LinalgExt::IREELinalgExtDialect,
+ IREE::Codegen::IREECodegenDialect>();
}
void runOnOperation() override;
};
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/lower_to_ukernel_ops.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/lower_to_ukernel_ops.mlir
index 20d7417..40be5f8 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/lower_to_ukernel_ops.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/lower_to_ukernel_ops.mlir
@@ -217,3 +217,18 @@
: tensor<?x?x7x8xf32> -> tensor<?x?xf32>
func.return %result : tensor<?x?xf32>
}
+
+// -----
+
+// CHECK: func @query_tile_sizes_2d(
+// CHECK-DAG: %[[DYNAMIC:.+]] = arith.constant -9223372036854775808 : index
+// CHECK-DAG: %[[FLAGS:.+]] = arith.constant 259 : i32
+// CHECK: %[[RESULT:.+]]:2 = iree_codegen.ukernel.generic "vmvx.query_tile_sizes.2d"
+// CHECK-SAME: ins(%[[DYNAMIC]], %[[DYNAMIC]], %[[FLAGS]] : index, index, i32)
+// CHECK: return %[[RESULT]]#0, %[[RESULT]]#1 : index, index
+func.func @query_tile_sizes_2d() -> (index, index) attributes {
+ hal.executable.target = #hal.executable.target<"vmvx", "vmvx-bytecode-fb">
+} {
+ %result:2 = iree_codegen.query_tile_sizes tensor<?x?xf32, #iree_linalg_ext.encoding<MATMUL_F32F32F32_RESULT>> -> index, index
+ return %result#0, %result#1 : index, index
+}
diff --git a/compiler/src/iree/compiler/Codegen/Transforms/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Transforms/BUILD.bazel
index 91d07a3..2ad5a2f 100644
--- a/compiler/src/iree/compiler/Codegen/Transforms/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Transforms/BUILD.bazel
@@ -24,6 +24,7 @@
],
deps = [
"//compiler/src/iree/compiler/Codegen/Utils",
+ "//compiler/src/iree/compiler/Codegen/Dialect:IREECodegenDialect",
"//compiler/src/iree/compiler/Dialect/Flow/IR",
"//compiler/src/iree/compiler/Dialect/HAL/IR",
# TODO(#13181) : Remove this dependency on VMVX dialect.
diff --git a/compiler/src/iree/compiler/Codegen/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Transforms/CMakeLists.txt
index c689604..199a950 100644
--- a/compiler/src/iree/compiler/Codegen/Transforms/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Transforms/CMakeLists.txt
@@ -35,6 +35,7 @@
MLIRTransforms
MLIRValueBoundsOpInterface
MLIRVectorDialect
+ iree::compiler::Codegen::Dialect::IREECodegenDialect
iree::compiler::Codegen::Utils
iree::compiler::Dialect::Flow::IR
iree::compiler::Dialect::HAL::IR
diff --git a/compiler/src/iree/compiler/Codegen/Transforms/Transforms.cpp b/compiler/src/iree/compiler/Codegen/Transforms/Transforms.cpp
index f76435f..f64351e 100644
--- a/compiler/src/iree/compiler/Codegen/Transforms/Transforms.cpp
+++ b/compiler/src/iree/compiler/Codegen/Transforms/Transforms.cpp
@@ -13,6 +13,7 @@
#include "iree/compiler/Codegen/Transforms/Transforms.h"
// TODO(#13038): Remove this dependency on VMVX dialect.
+#include "iree/compiler/Codegen/Dialect/IREECodegenOps.h"
#include "iree/compiler/Dialect/VMVX/IR/VMVXOps.h"
#include "llvm/Support/Debug.h"
#include "mlir/Analysis/Liveness.h"
@@ -311,7 +312,7 @@
// TODO(#13038) This is a WAR for the these ops ending up in workgroup count
// computation. They should not. Some pre-processing at MaterializeEncoding
// time might make these go away.
- if (isa<IREE::VMVX::QueryTileSizesOp>(op)) {
+ if (isa<IREE::Codegen::QueryTileSizesOp>(op)) {
Value constVal =
rewriter.create<arith::ConstantIndexOp>(op->getLoc(), 16);
for (auto result : op->getResults()) {
diff --git a/compiler/src/iree/compiler/Codegen/Utils/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Utils/BUILD.bazel
index ea4a2ad..db2f2b3 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Utils/BUILD.bazel
@@ -17,7 +17,7 @@
iree_compiler_cc_library(
name = "Utils",
srcs = [
- "EncodingInfo.cpp",
+ "EncodingUtils.cpp",
"GPUUtils.cpp",
"LinalgOpInfo.cpp",
"LinkingUtils.cpp",
@@ -25,7 +25,7 @@
"Utils.cpp",
],
hdrs = [
- "EncodingInfo.h",
+ "EncodingUtils.h",
"GPUUtils.h",
"LinalgOpInfo.h",
"LinkingUtils.h",
diff --git a/compiler/src/iree/compiler/Codegen/Utils/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Utils/CMakeLists.txt
index 6c6a1bf..863d703 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Utils/CMakeLists.txt
@@ -14,14 +14,14 @@
NAME
Utils
HDRS
- "EncodingInfo.h"
+ "EncodingUtils.h"
"GPUUtils.h"
"LinalgOpInfo.h"
"LinkingUtils.h"
"MarkerUtils.h"
"Utils.h"
SRCS
- "EncodingInfo.cpp"
+ "EncodingUtils.cpp"
"GPUUtils.cpp"
"LinalgOpInfo.cpp"
"LinkingUtils.cpp"
diff --git a/compiler/src/iree/compiler/Codegen/Utils/EncodingInfo.cpp b/compiler/src/iree/compiler/Codegen/Utils/EncodingInfo.cpp
deleted file mode 100644
index 10cc0e1..0000000
--- a/compiler/src/iree/compiler/Codegen/Utils/EncodingInfo.cpp
+++ /dev/null
@@ -1,32 +0,0 @@
-// Copyright 2023 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/Utils/EncodingInfo.h"
-
-#include "mlir/IR/BuiltinTypes.h"
-
-namespace mlir {
-namespace iree_compiler {
-
-std::optional<MatmulType> getMatmulType(Type lhsElementType,
- Type rhsElementType,
- Type resultElementType) {
- if (lhsElementType.isSignlessInteger(8) &&
- rhsElementType.isSignlessInteger(8) &&
- resultElementType.isSignlessInteger(32)) {
- return MatmulType::I8I8I32;
- }
-
- if (lhsElementType.isF32() && rhsElementType.isF32() &&
- resultElementType.isF32()) {
- return MatmulType::F32F32F32;
- }
-
- return std::nullopt;
-}
-
-} // namespace iree_compiler
-} // namespace mlir
diff --git a/compiler/src/iree/compiler/Codegen/Utils/EncodingInfo.h b/compiler/src/iree/compiler/Codegen/Utils/EncodingInfo.h
deleted file mode 100644
index d40b3f4..0000000
--- a/compiler/src/iree/compiler/Codegen/Utils/EncodingInfo.h
+++ /dev/null
@@ -1,28 +0,0 @@
-// Copyright 2023 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_SRC_IREE_COMPILER_CODEGEN_UTILS_ENCODINGINFO_H_
-#define IREE_COMPILER_SRC_IREE_COMPILER_CODEGEN_UTILS_ENCODINGINFO_H_
-
-#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h"
-#include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
-
-namespace mlir {
-namespace iree_compiler {
-
-enum class MatmulType {
- F32F32F32,
- I8I8I32,
-};
-
-std::optional<MatmulType> getMatmulType(Type lhsElementType,
- Type rhsElementType,
- Type resultElementType);
-
-} // namespace iree_compiler
-} // namespace mlir
-
-#endif // IREE_COMPILER_SRC_IREE_COMPILER_CODEGEN_UTILS_ENCODINGINFO_H_
diff --git a/compiler/src/iree/compiler/Codegen/Utils/EncodingUtils.cpp b/compiler/src/iree/compiler/Codegen/Utils/EncodingUtils.cpp
new file mode 100644
index 0000000..8ba4384
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Utils/EncodingUtils.cpp
@@ -0,0 +1,74 @@
+// Copyright 2023 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/Utils/EncodingUtils.h"
+
+#include "mlir/IR/BuiltinTypes.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+using IREE::LinalgExt::EncodingAttr;
+using IREE::LinalgExt::TensorEncoding;
+using IREE::LinalgExt::TensorEncodingAttr;
+
+std::optional<MatmulType> getMatmulType(Type lhsElementType,
+ Type rhsElementType,
+ Type resultElementType) {
+ if (lhsElementType.isSignlessInteger(8) &&
+ rhsElementType.isSignlessInteger(8) &&
+ resultElementType.isSignlessInteger(32)) {
+ return MatmulType::I8I8I32;
+ }
+
+ if (lhsElementType.isF32() && rhsElementType.isF32() &&
+ resultElementType.isF32()) {
+ return MatmulType::F32F32F32;
+ }
+
+ return std::nullopt;
+}
+
+std::optional<TensorEncoding> getEncoding(RankedTensorType tensorType) {
+ auto encodingAttr =
+ llvm::dyn_cast_if_present<EncodingAttr>(tensorType.getEncoding());
+ if (!encodingAttr) return std::nullopt;
+ return encodingAttr.getEncoding().getValue();
+}
+
+std::optional<MatmulType> getMatmulType(TensorEncoding encoding) {
+ switch (encoding) {
+ case TensorEncoding::MATMUL_F32F32F32_LHS:
+ case TensorEncoding::MATMUL_F32F32F32_RHS:
+ case TensorEncoding::MATMUL_F32F32F32_RESULT:
+ return MatmulType::F32F32F32;
+ case TensorEncoding::MATMUL_I8I8I32_LHS:
+ case TensorEncoding::MATMUL_I8I8I32_RHS:
+ case TensorEncoding::MATMUL_I8I8I32_RESULT:
+ return MatmulType::I8I8I32;
+ default:
+ return std::nullopt;
+ }
+}
+
+std::optional<MatmulOperandRole> getMatmulOperandRole(TensorEncoding encoding) {
+ switch (encoding) {
+ case TensorEncoding::MATMUL_F32F32F32_LHS:
+ case TensorEncoding::MATMUL_I8I8I32_LHS:
+ return MatmulOperandRole::LHS;
+ case TensorEncoding::MATMUL_F32F32F32_RHS:
+ case TensorEncoding::MATMUL_I8I8I32_RHS:
+ return MatmulOperandRole::RHS;
+ case TensorEncoding::MATMUL_F32F32F32_RESULT:
+ case TensorEncoding::MATMUL_I8I8I32_RESULT:
+ return MatmulOperandRole::RESULT;
+ default:
+ return std::nullopt;
+ }
+}
+
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/compiler/src/iree/compiler/Codegen/Utils/EncodingUtils.h b/compiler/src/iree/compiler/Codegen/Utils/EncodingUtils.h
new file mode 100644
index 0000000..c6a2100
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Utils/EncodingUtils.h
@@ -0,0 +1,54 @@
+// Copyright 2023 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_SRC_IREE_COMPILER_CODEGEN_UTILS_ENCODINGUTILS_H_
+#define IREE_COMPILER_SRC_IREE_COMPILER_CODEGEN_UTILS_ENCODINGUTILS_H_
+
+#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+// Enumeration of possible (LHS, RHS, Accumulator) type triples for the
+// element types of the operands of a matmul-like. or linalg.mmt4d.
+enum class MatmulType {
+ F32F32F32,
+ I8I8I32,
+};
+
+// Enumeration of the operands of a matmul-like operation such as linalg.matmul.
+enum class MatmulOperandRole {
+ LHS,
+ RHS,
+ RESULT,
+};
+
+// Constructs a MatmulType from separate operands element types, or returns
+// std::nullopt if no MatmulType enumeration value would match.
+std::optional<MatmulType> getMatmulType(Type lhsElementType,
+ Type rhsElementType,
+ Type resultElementType);
+
+// Helper to read the TensorEncoding from a TensorEncodingAttr on a TensorType.
+// Return std::nullopt if the TensorType does not have a TensorEncodingAttr.
+std::optional<IREE::LinalgExt::TensorEncoding> getEncoding(
+ RankedTensorType tensorType);
+
+// Reads a MatmulType from a TensorEncoding, or returns std::nullopt if no
+// MatmulType enumeration value would match.
+std::optional<MatmulType> getMatmulType(
+ IREE::LinalgExt::TensorEncoding encoding);
+
+// Reads a MatmulOperandRole from a TensorEncoding, or returns std::nullopt if
+// no MatmulOperandRole enumeration value would match.
+std::optional<MatmulOperandRole> getMatmulOperandRole(
+ IREE::LinalgExt::TensorEncoding encoding);
+
+} // namespace iree_compiler
+} // namespace mlir
+
+#endif // IREE_COMPILER_SRC_IREE_COMPILER_CODEGEN_UTILS_ENCODINGUTILS_H_
diff --git a/compiler/src/iree/compiler/Codegen/VMVX/BUILD.bazel b/compiler/src/iree/compiler/Codegen/VMVX/BUILD.bazel
index b93deb0..fcc30b9 100644
--- a/compiler/src/iree/compiler/Codegen/VMVX/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/VMVX/BUILD.bazel
@@ -29,6 +29,7 @@
"//compiler/src/iree/compiler/Codegen:PassHeaders",
"//compiler/src/iree/compiler/Codegen/Common",
"//compiler/src/iree/compiler/Codegen/Common:CommonPasses",
+ "//compiler/src/iree/compiler/Codegen/Dialect:IREECodegenDialect",
"//compiler/src/iree/compiler/Codegen/Utils",
"//compiler/src/iree/compiler/Dialect/Flow/IR",
"//compiler/src/iree/compiler/Dialect/HAL/IR",
diff --git a/compiler/src/iree/compiler/Codegen/VMVX/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/VMVX/CMakeLists.txt
index a337a81..d9dea0d 100644
--- a/compiler/src/iree/compiler/Codegen/VMVX/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/VMVX/CMakeLists.txt
@@ -41,6 +41,7 @@
iree::builtins::ukernel::exported_bits
iree::compiler::Codegen::Common
iree::compiler::Codegen::Common::CommonPasses
+ iree::compiler::Codegen::Dialect::IREECodegenDialect
iree::compiler::Codegen::PassHeaders
iree::compiler::Codegen::Utils
iree::compiler::Dialect::Flow::IR
diff --git a/compiler/src/iree/compiler/Codegen/VMVX/VMVXMaterializeEncodingPass.cpp b/compiler/src/iree/compiler/Codegen/VMVX/VMVXMaterializeEncodingPass.cpp
index d83ebdb..e2f052d 100644
--- a/compiler/src/iree/compiler/Codegen/VMVX/VMVXMaterializeEncodingPass.cpp
+++ b/compiler/src/iree/compiler/Codegen/VMVX/VMVXMaterializeEncodingPass.cpp
@@ -7,6 +7,7 @@
#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtDialect.h"
#include "iree-dialects/Dialect/LinalgExt/Passes/Passes.h"
#include "iree/compiler/Codegen/Common/EncodingInfo.h"
+#include "iree/compiler/Codegen/Dialect/IREECodegenDialect.h"
#include "iree/compiler/Codegen/PassDetail.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "iree/compiler/Codegen/VMVX/EncodingInfo.h"
@@ -55,10 +56,10 @@
struct VMVXMaterializeEncodingPass
: public VMVXMaterializeEncodingBase<VMVXMaterializeEncodingPass> {
void getDependentDialects(DialectRegistry ®istry) const override {
- registry.insert<arith::ArithDialect, affine::AffineDialect,
- tensor::TensorDialect, IREE::Flow::FlowDialect,
- IREE::LinalgExt::IREELinalgExtDialect,
- IREE::VMVX::VMVXDialect>();
+ registry.insert<
+ arith::ArithDialect, affine::AffineDialect, tensor::TensorDialect,
+ IREE::Flow::FlowDialect, IREE::LinalgExt::IREELinalgExtDialect,
+ IREE::VMVX::VMVXDialect, IREE::Codegen::IREECodegenDialect>();
}
void runOnOperation() override;
};
diff --git a/compiler/src/iree/compiler/Codegen/VMVX/test/materialize_encoding.mlir b/compiler/src/iree/compiler/Codegen/VMVX/test/materialize_encoding.mlir
index bab503d..1446ebc 100644
--- a/compiler/src/iree/compiler/Codegen/VMVX/test/materialize_encoding.mlir
+++ b/compiler/src/iree/compiler/Codegen/VMVX/test/materialize_encoding.mlir
@@ -36,21 +36,20 @@
// CHECK-DAG: #[[MAP_CEILDIV:.+]] = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
// CHECK: func @matmul_lowering_i8i8i32_vmvx_ukernel()
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
-// CHECK-DAG: %[[DYNAMIC:.+]] = arith.constant -9223372036854775808 : index
// CHECK-DAG: %[[M:.+]] = hal.interface.constant.load[0]
// CHECK-DAG: %[[N:.+]] = hal.interface.constant.load[1]
// CHECK-DAG: %[[K:.+]] = hal.interface.constant.load[2]
-// CHECK: %[[LHS_TILE_SIZES:.+]]:2 = vmvx.query_tile_sizes sizes(%[[DYNAMIC]], %[[DYNAMIC]]) flags(513) -> index, index
+// CHECK: %[[LHS_TILE_SIZES:.+]]:2 = iree_codegen.query_tile_sizes tensor<?x?xi8, #iree_linalg_ext.encoding<MATMUL_I8I8I32_LHS>> -> index, index
// CHECK-DAG: %[[LHS_OUTER_SIZE0:.+]] = affine.apply #[[MAP_CEILDIV]]()[%[[M]], %[[LHS_TILE_SIZES]]#0]
// CHECK-DAG: %[[LHS_OUTER_SIZE1:.+]] = affine.apply #[[MAP_CEILDIV]]()[%[[K]], %[[LHS_TILE_SIZES]]#1]
// CHECK: %[[LHS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(0)
// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi8>>{%[[LHS_OUTER_SIZE0]], %[[LHS_OUTER_SIZE1]], %[[LHS_TILE_SIZES]]#0, %[[LHS_TILE_SIZES]]#1}
-// CHECK: %[[RHS_TILE_SIZES:.+]]:2 = vmvx.query_tile_sizes sizes(%[[DYNAMIC]], %[[DYNAMIC]]) flags(514) -> index, index
+// CHECK: %[[RHS_TILE_SIZES:.+]]:2 = iree_codegen.query_tile_sizes tensor<?x?xi8, #iree_linalg_ext.encoding<MATMUL_I8I8I32_RHS>> -> index, index
// CHECK-DAG: %[[RHS_OUTER_SIZE0:.+]] = affine.apply #[[MAP_CEILDIV]]()[%[[N]], %[[RHS_TILE_SIZES]]#0]
// CHECK-DAG: %[[RHS_OUTER_SIZE1:.+]] = affine.apply #[[MAP_CEILDIV]]()[%[[K]], %[[RHS_TILE_SIZES]]#1]
// CHECK: %[[RHS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(1)
// CHECK-SAME: !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi8>>{%[[RHS_OUTER_SIZE0]], %[[RHS_OUTER_SIZE1]], %[[RHS_TILE_SIZES]]#0, %[[RHS_TILE_SIZES]]#1}
-// CHECK: %[[RESULT_TILE_SIZES:.+]]:2 = vmvx.query_tile_sizes sizes(%[[DYNAMIC]], %[[DYNAMIC]]) flags(515) -> index, index
+// CHECK: %[[RESULT_TILE_SIZES:.+]]:2 = iree_codegen.query_tile_sizes tensor<?x?xi32, #iree_linalg_ext.encoding<MATMUL_I8I8I32_RESULT>> -> index, index
// CHECK-DAG: %[[RESULT_OUTER_SIZE0:.+]] = affine.apply #[[MAP_CEILDIV]]()[%[[M]], %[[RESULT_TILE_SIZES]]#0]
// CHECK-DAG: %[[RESULT_OUTER_SIZE1:.+]] = affine.apply #[[MAP_CEILDIV]]()[%[[N]], %[[RESULT_TILE_SIZES]]#1]
// CHECK: %[[OUTS_BINDING:.+]] = hal.interface.binding.subspan set(0) binding(2)
diff --git a/compiler/src/iree/compiler/Codegen/VMVX/test/pipeline.mlir b/compiler/src/iree/compiler/Codegen/VMVX/test/pipeline.mlir
index e06bbf8..22b0e64 100644
--- a/compiler/src/iree/compiler/Codegen/VMVX/test/pipeline.mlir
+++ b/compiler/src/iree/compiler/Codegen/VMVX/test/pipeline.mlir
@@ -13,35 +13,23 @@
%c256 = arith.constant 256 : index
%c512 = arith.constant 512 : index
%c16 = arith.constant 16 : index
- %0:2 = vmvx.query_tile_sizes sizes(%c16, %c16) flags(1048576) -> index, index
+ %0:2 = iree_codegen.query_tile_sizes tensor<16x16xi8, #iree_linalg_ext.encoding<MATMUL_I8I8I32_LHS>> -> index, index
%1 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%0#0]
%2 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%0#1]
%3 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi8>>{%1, %2, %0#0, %0#1}
- %4:2 = vmvx.query_tile_sizes sizes(%c16, %c16) flags(1114112) -> index, index
+ %4:2 = iree_codegen.query_tile_sizes tensor<16x16xi8, #iree_linalg_ext.encoding<MATMUL_I8I8I32_RHS>> -> index, index
%5 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%4#0]
%6 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%4#1]
%7 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c256) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi8>>{%5, %6, %4#0, %4#1}
- %8:2 = vmvx.query_tile_sizes sizes(%c16, %c16) flags(1179648) -> index, index
+ %8:2 = iree_codegen.query_tile_sizes tensor<16x16xi32, #iree_linalg_ext.encoding<MATMUL_I8I8I32_RESULT>> -> index, index
%9 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%8#0]
%10 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%8#1]
%11 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c512) : !flow.dispatch.tensor<readwrite:tensor<?x?x?x?xi32>>{%9, %10, %8#0, %8#1}
- %12:2 = vmvx.query_tile_sizes sizes(%c16, %c16) flags(1048576) -> index, index
- %13 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%12#0]
- %14 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%12#1]
- %15 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0], sizes = [%13, %14, %12#0, %12#1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi8>>{%13, %14, %12#0, %12#1} -> tensor<?x?x?x?xi8>
- %16:2 = vmvx.query_tile_sizes sizes(%c16, %c16) flags(1114112) -> index, index
- %17 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%16#0]
- %18 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%16#1]
- %19 = flow.dispatch.tensor.load %7, offsets = [0, 0, 0, 0], sizes = [%17, %18, %16#0, %16#1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi8>>{%17, %18, %16#0, %16#1} -> tensor<?x?x?x?xi8>
- %20:2 = vmvx.query_tile_sizes sizes(%c16, %c16) flags(1179648) -> index, index
- %21 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%20#0]
- %22 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%20#1]
- %23 = flow.dispatch.tensor.load %11, offsets = [0, 0, 0, 0], sizes = [%21, %22, %20#0, %20#1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readwrite:tensor<?x?x?x?xi32>>{%21, %22, %20#0, %20#1} -> tensor<?x?x?x?xi32>
+ %15 = flow.dispatch.tensor.load %3, offsets = [0, 0, 0, 0], sizes = [%1, %2, %0#0, %0#1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi8>>{%1, %2, %0#0, %0#1} -> tensor<?x?x?x?xi8>
+ %19 = flow.dispatch.tensor.load %7, offsets = [0, 0, 0, 0], sizes = [%5, %6, %4#0, %4#1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi8>>{%5, %6, %4#0, %4#1} -> tensor<?x?x?x?xi8>
+ %23 = flow.dispatch.tensor.load %11, offsets = [0, 0, 0, 0], sizes = [%9, %10, %8#0, %8#1], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readwrite:tensor<?x?x?x?xi32>>{%9, %10, %8#0, %8#1} -> tensor<?x?x?x?xi32>
%24 = linalg.mmt4d ins(%15, %19 : tensor<?x?x?x?xi8>, tensor<?x?x?x?xi8>) outs(%23 : tensor<?x?x?x?xi32>) -> tensor<?x?x?x?xi32>
- %25:2 = vmvx.query_tile_sizes sizes(%c16, %c16) flags(1179648) -> index, index
- %26 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%25#0]
- %27 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%25#1]
- flow.dispatch.tensor.store %24, %11, offsets = [0, 0, 0, 0], sizes = [%26, %27, %25#0, %25#1], strides = [1, 1, 1, 1] : tensor<?x?x?x?xi32> -> !flow.dispatch.tensor<readwrite:tensor<?x?x?x?xi32>>{%26, %27, %25#0, %25#1}
+ flow.dispatch.tensor.store %24, %11, offsets = [0, 0, 0, 0], sizes = [%9, %10, %8#0, %8#1], strides = [1, 1, 1, 1] : tensor<?x?x?x?xi32> -> !flow.dispatch.tensor<readwrite:tensor<?x?x?x?xi32>>{%9, %10, %8#0, %8#1}
return
}
}
diff --git a/compiler/src/iree/compiler/Dialect/VMVX/Conversion/VMVXToVM/ConvertVMVXToVM.cpp b/compiler/src/iree/compiler/Dialect/VMVX/Conversion/VMVXToVM/ConvertVMVXToVM.cpp
index 7097b65..2e32ed9 100644
--- a/compiler/src/iree/compiler/Dialect/VMVX/Conversion/VMVXToVM/ConvertVMVXToVM.cpp
+++ b/compiler/src/iree/compiler/Dialect/VMVX/Conversion/VMVXToVM/ConvertVMVXToVM.cpp
@@ -160,17 +160,6 @@
}
};
-// Converts the vmvx.query_tile_sizes op to its import.
-class QueryTileSizesOpConversion
- : public VMVXImportOpConversion<IREE::VMVX::QueryTileSizesOp> {
- public:
- using VMVXImportOpConversion::VMVXImportOpConversion;
-
- std::string getImportFqName(IREE::VMVX::QueryTileSizesOp op) const override {
- return "vmvx.query_tile_sizes.2d";
- }
-};
-
class UnaryOpConversion : public VMVXImportOpConversion<IREE::VMVX::UnaryOp> {
public:
using VMVXImportOpConversion::VMVXImportOpConversion;
@@ -195,8 +184,7 @@
SymbolTable &importSymbols,
RewritePatternSet &patterns) {
patterns.insert<BinaryOpConversion, CopyOpConversion, Fill2DOpConversion,
- UnaryOpConversion, QueryTileSizesOpConversion>(
- context, importSymbols, typeConverter);
+ UnaryOpConversion>(context, importSymbols, typeConverter);
}
} // namespace iree_compiler
diff --git a/compiler/src/iree/compiler/Dialect/VMVX/Conversion/VMVXToVM/test/BUILD.bazel b/compiler/src/iree/compiler/Dialect/VMVX/Conversion/VMVXToVM/test/BUILD.bazel
index 994375f..e90b2b5 100644
--- a/compiler/src/iree/compiler/Dialect/VMVX/Conversion/VMVXToVM/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Dialect/VMVX/Conversion/VMVXToVM/test/BUILD.bazel
@@ -19,7 +19,6 @@
"binary.mlir",
"copy.mlir",
"fill.mlir",
- "query_tile_sizes.mlir",
"unary.mlir",
],
include = ["*.mlir"],
diff --git a/compiler/src/iree/compiler/Dialect/VMVX/Conversion/VMVXToVM/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/VMVX/Conversion/VMVXToVM/test/CMakeLists.txt
index df49dcb..0aae7e7 100644
--- a/compiler/src/iree/compiler/Dialect/VMVX/Conversion/VMVXToVM/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/VMVX/Conversion/VMVXToVM/test/CMakeLists.txt
@@ -17,7 +17,6 @@
"binary.mlir"
"copy.mlir"
"fill.mlir"
- "query_tile_sizes.mlir"
"unary.mlir"
TOOLS
FileCheck
diff --git a/compiler/src/iree/compiler/Dialect/VMVX/Conversion/VMVXToVM/test/query_tile_sizes.mlir b/compiler/src/iree/compiler/Dialect/VMVX/Conversion/VMVXToVM/test/query_tile_sizes.mlir
deleted file mode 100644
index 0b65b02..0000000
--- a/compiler/src/iree/compiler/Dialect/VMVX/Conversion/VMVXToVM/test/query_tile_sizes.mlir
+++ /dev/null
@@ -1,13 +0,0 @@
-// RUN: iree-opt --iree-vm-target-index-bits=64 --split-input-file \
-// RUN: --iree-vm-conversion --canonicalize %s | FileCheck %s
-
-// CHECK-LABEL: @foo
-func.func @foo(
- // Sizes
- %arg0 : index, %arg1 : index
- ) -> (index, index) {
- // CHECK: %[[FLAGS:.+]] = vm.const.i32 1048576
- // CHECK: vm.call @vmvx.query_tile_sizes.2d(%arg0, %arg1, %[[FLAGS]]) : (i64, i64, i32) -> (i64, i64)
- %tile_sizes:2 = vmvx.query_tile_sizes sizes (%arg0, %arg1) flags (1048576) -> index, index
- return %tile_sizes#0, %tile_sizes#1 : index, index
-}
diff --git a/compiler/src/iree/compiler/Dialect/VMVX/IR/VMVXOps.td b/compiler/src/iree/compiler/Dialect/VMVX/IR/VMVXOps.td
index 4e4d13f..8a91550 100644
--- a/compiler/src/iree/compiler/Dialect/VMVX/IR/VMVXOps.td
+++ b/compiler/src/iree/compiler/Dialect/VMVX/IR/VMVXOps.td
@@ -75,27 +75,6 @@
}];
}
-def VMVX_QueryTileSizesOp : VMVX_PureOp<"query_tile_sizes", [Pure]> {
- let summary = "Returns tile sizes to use to materialize the given encoding, based on runtime CPU capabilities.";
- let description = [{
- TODO write me
- }];
-
- let arguments = (ins
- Variadic<VMVX_Index>:$sizes,
- I32Attr:$flags
- );
- let results = (outs
- Variadic<VMVX_Index>:$tileSizes
- );
-
- let assemblyFormat = [{
- `sizes` `` `(` $sizes `)`
- `flags` `` `(` $flags `)`
- `->` type(results) attr-dict
- }];
-}
-
//===----------------------------------------------------------------------===//
// VMVX Ops: ABI
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Dialect/VMVX/vmvx.imports.mlir b/compiler/src/iree/compiler/Dialect/VMVX/vmvx.imports.mlir
index 2ab7438..fc03b09 100644
--- a/compiler/src/iree/compiler/Dialect/VMVX/vmvx.imports.mlir
+++ b/compiler/src/iree/compiler/Dialect/VMVX/vmvx.imports.mlir
@@ -482,7 +482,8 @@
//==============================================================================
vm.import private @query_tile_sizes.2d(
- %sizes : tuple<i64, i64>,
+ %size0 : i64,
+ %size1 : i64,
%flags : i32
) -> (i64, i64)
