From 1f9bf3a830883ef5af6796383b499e269ac46ea6 Mon Sep 17 00:00:00 2001 From: Han-Chung Wang Date: Thu, 21 Nov 2024 15:47:50 -0800 Subject: [PATCH] [Codegen][NFC] Move encoding related containers to Dialect/Codegen/Utils (#19234) The revision moves MaterializeEncodingInfo struct and TileSwizzle struct to `compiler/Codegen/Dialect/Codegen/Utils/Utils.[h|cpp]`. It is a preparation for https://github.com/iree-org/iree/issues/17924 because they are how we define layouts in data-tiling and we're going to expose the layouts to EncodingAttr. The revision also updates the namespace in ` GPUTileSwizzleUtils.[h|cpp]`, which follows the convention; a typo in license. They are all created in 2024, so the year should be 2024: https://github.com/iree-org/iree/commit/740e301d61e18a5833e1a2d75b476ae850f8c17e --------- Signed-off-by: hanhanW --- .../iree/compiler/Codegen/Common/BUILD.bazel | 9 +---- .../compiler/Codegen/Common/CMakeLists.txt | 14 +------ .../compiler/Codegen/Common/CPU/BUILD.bazel | 1 + .../Codegen/Common/CPU/CMakeLists.txt | 1 + .../Common/CPU/CPUMaterializeEncodings.cpp | 2 + .../compiler/Codegen/Common/EncodingUtils.cpp | 15 ++------ .../compiler/Codegen/Common/EncodingUtils.h | 28 +++----------- .../compiler/Codegen/Common/GPU/BUILD.bazel | 1 + .../Codegen/Common/GPU/CMakeLists.txt | 1 + .../Common/GPU/GPUMaterializeEncoding.cpp | 4 ++ .../Common/MaterializeEncodingIntoNop.cpp | 1 + .../MaterializeEncodingIntoPackUnPack.cpp | 4 +- .../Codegen/Dialect/Codegen/Utils/BUILD.bazel | 28 ++++++++++++++ .../Dialect/Codegen/Utils/CMakeLists.txt | 27 +++++++++++++ .../Codegen/Utils/Utils.cpp} | 25 ++++++++++-- .../Codegen/Utils/Utils.h} | 38 ++++++++++++++++--- .../Codegen/Dialect/GPU/IR/BUILD.bazel | 2 +- .../Codegen/Dialect/GPU/IR/CMakeLists.txt | 2 +- .../Dialect/GPU/IR/GPUTileSwizzleUtils.cpp | 8 +++- .../Dialect/GPU/IR/GPUTileSwizzleUtils.h | 16 ++++---- .../Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp | 2 + .../LLVMGPU/TransformExtensions/BUILD.bazel | 1 + .../TransformExtensions/CMakeLists.txt | 1 + .../TransformExtensions/LLVMGPUExtensions.cpp | 2 + 24 files changed, 157 insertions(+), 76 deletions(-) create mode 100644 compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/BUILD.bazel create mode 100644 compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/CMakeLists.txt rename compiler/src/iree/compiler/Codegen/{Common/TileSwizzle.cpp => Dialect/Codegen/Utils/Utils.cpp} (63%) rename compiler/src/iree/compiler/Codegen/{Common/TileSwizzle.h => Dialect/Codegen/Utils/Utils.h} (72%) diff --git a/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel index ee80c895c805..ece9b792408a 100644 --- a/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel +++ b/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel @@ -171,8 +171,8 @@ iree_compiler_cc_library( ":PassHeaders", ":PassesIncGen", "//compiler/src/iree/compiler/Codegen/Common:FoldTensorExtractOpIncGen", - "//compiler/src/iree/compiler/Codegen/Common:TileSwizzle", "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect", + "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils", "//compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR:IREEVectorExtDialect", "//compiler/src/iree/compiler/Codegen/Interfaces:BufferizationInterfaces", "//compiler/src/iree/compiler/Codegen/Interfaces:PartitionableLoopsInterface", @@ -240,13 +240,6 @@ iree_compiler_cc_library( ], ) -iree_compiler_cc_library( - name = "TileSwizzle", - srcs = ["TileSwizzle.cpp"], - hdrs = ["TileSwizzle.h"], - deps = ["@llvm-project//llvm:Support"], -) - # TODO: If the layering causes concerns then the transform dialect interpreter # should be one level above everything: it is a mechanism by which # transformations are applied to any IR and needs to register all the dialects diff --git a/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt index 9ce0c7ea8fd9..9a39a10b81a3 100644 --- a/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt +++ b/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt @@ -199,8 +199,8 @@ iree_cc_library( MLIRVectorTransforms MLIRViewLikeInterface iree::compiler::Codegen::Common::FoldTensorExtractOpIncGen - iree::compiler::Codegen::Common::TileSwizzle iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect + iree::compiler::Codegen::Dialect::Codegen::Utils iree::compiler::Codegen::Dialect::VectorExt::IR::IREEVectorExtDialect iree::compiler::Codegen::Interfaces::BufferizationInterfaces iree::compiler::Codegen::Interfaces::PartitionableLoopsInterface @@ -218,18 +218,6 @@ iree_cc_library( PUBLIC ) -iree_cc_library( - NAME - TileSwizzle - HDRS - "TileSwizzle.h" - SRCS - "TileSwizzle.cpp" - DEPS - LLVMSupport - PUBLIC -) - iree_cc_library( NAME TransformDialectInterpreterPass diff --git a/compiler/src/iree/compiler/Codegen/Common/CPU/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/CPU/BUILD.bazel index acac444d2780..d6b3c96d1108 100644 --- a/compiler/src/iree/compiler/Codegen/Common/CPU/BUILD.bazel +++ b/compiler/src/iree/compiler/Codegen/Common/CPU/BUILD.bazel @@ -57,6 +57,7 @@ iree_compiler_cc_library( ":PassesIncGen", "//compiler/src/iree/compiler/Codegen/Common", "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect", + "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils", "//compiler/src/iree/compiler/Codegen/Interfaces:UKernelOpInterface", "//compiler/src/iree/compiler/Codegen/Transforms", "//compiler/src/iree/compiler/Codegen/Utils", diff --git a/compiler/src/iree/compiler/Codegen/Common/CPU/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/CPU/CMakeLists.txt index 8b20fdfc321e..81cefe71999f 100644 --- a/compiler/src/iree/compiler/Codegen/Common/CPU/CMakeLists.txt +++ b/compiler/src/iree/compiler/Codegen/Common/CPU/CMakeLists.txt @@ -79,6 +79,7 @@ iree_cc_library( iree::builtins::ukernel::exported_bits iree::compiler::Codegen::Common iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect + iree::compiler::Codegen::Dialect::Codegen::Utils iree::compiler::Codegen::Interfaces::UKernelOpInterface iree::compiler::Codegen::Transforms iree::compiler::Codegen::Utils diff --git a/compiler/src/iree/compiler/Codegen/Common/CPU/CPUMaterializeEncodings.cpp b/compiler/src/iree/compiler/Codegen/Common/CPU/CPUMaterializeEncodings.cpp index 2ccc25ce4f7d..18fabd49a4a5 100644 --- a/compiler/src/iree/compiler/Codegen/Common/CPU/CPUMaterializeEncodings.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/CPU/CPUMaterializeEncodings.cpp @@ -30,6 +30,8 @@ namespace mlir::iree_compiler { +using IREE::Codegen::MaterializeEncodingInfo; + #define GEN_PASS_DEF_CPUMATERIALIZEDEVICEENCODINGPASS #define GEN_PASS_DEF_CPUMATERIALIZEHOSTENCODINGPASS #include "iree/compiler/Codegen/Common/CPU/Passes.h.inc" diff --git a/compiler/src/iree/compiler/Codegen/Common/EncodingUtils.cpp b/compiler/src/iree/compiler/Codegen/Common/EncodingUtils.cpp index e71960eec778..7d041d09a738 100644 --- a/compiler/src/iree/compiler/Codegen/Common/EncodingUtils.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/EncodingUtils.cpp @@ -5,6 +5,7 @@ // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception #include "iree/compiler/Codegen/Common/EncodingUtils.h" +#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h" #include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Utils/IndexingUtils.h" @@ -14,6 +15,7 @@ namespace mlir::iree_compiler { +using IREE::Codegen::MaterializeEncodingInfo; using IREE::Encoding::EncodingAttr; using IREE::Encoding::getEncodingAttr; using IREE::Encoding::getEncodingContractionDims; @@ -120,7 +122,7 @@ MaterializeEncodingTypeConverter::MaterializeEncodingTypeConverter( SmallVector newShape( packedType.getShape().drop_back(encodingInfo.innerTileSizes.size())); SmallVector swizzledTileShape = - getExpandedTileShape(swizzle.expandShape); + IREE::Codegen::getExpandedTileShape(swizzle.expandShape); applyPermutationToVector(swizzledTileShape, swizzle.permutation); newShape.append(swizzledTileShape); return RankedTensorType::get(newShape, packedType.getElementType()); @@ -199,15 +201,4 @@ bool isNarrowNResult(EncodingAttr encoding) { return IREE::Encoding::getMatmulNarrowDim(encoding).isN(); } -SmallVector -getExpandedTileShape(const TileSwizzle::ExpandShapeType &expandShape) { - SmallVector result; - for (auto e : expandShape) { - for (auto d : e) { - result.push_back(d.size); - } - } - return result; -} - } // namespace mlir::iree_compiler diff --git a/compiler/src/iree/compiler/Codegen/Common/EncodingUtils.h b/compiler/src/iree/compiler/Codegen/Common/EncodingUtils.h index b30ef24b87ef..1c9d0860c5d8 100644 --- a/compiler/src/iree/compiler/Codegen/Common/EncodingUtils.h +++ b/compiler/src/iree/compiler/Codegen/Common/EncodingUtils.h @@ -7,7 +7,7 @@ #ifndef IREE_COMPILER_SRC_IREE_COMPILER_CODEGEN_COMMON_ENCODINGUTILS_H_ #define IREE_COMPILER_SRC_IREE_COMPILER_CODEGEN_COMMON_ENCODINGUTILS_H_ -#include "iree/compiler/Codegen/Common/TileSwizzle.h" +#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h" #include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h" #include "iree/compiler/Dialect/HAL/IR/HALTypes.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" @@ -15,21 +15,9 @@ namespace mlir::iree_compiler { -/// Container of information needed to materialize the layout transformations. -struct MaterializeEncodingInfo { - // The next 3 fields are used to create a `tensor.pack` or `tensor.unpack` op, - // changing the overall layout between row-major and tiled (where each tile is - // row-major). - SmallVector innerDimsPos; - SmallVector innerTileSizes; - SmallVector outerDimsPerm; - - // The optional swizzle, see the comment on TileSwizzle. Only used on GPU. - std::optional swizzle; -}; - -using MaterializeEncodingFn = std::function( - RankedTensorType, IREE::HAL::ExecutableTargetAttr targetAttr)>; +using MaterializeEncodingFn = + std::function( + RankedTensorType, IREE::HAL::ExecutableTargetAttr targetAttr)>; struct MaterializeEncodingValueInfo { SmallVector innerTileSizes; @@ -56,7 +44,7 @@ class MaterializeEncodingTypeConverter : public TypeConverter { IREE::HAL::ExecutableTargetAttr getTargetAttr() const { return targetAttr; } - FailureOr + FailureOr getEncodingInfo(RankedTensorType type) const { return materializeEncodingFn(type, targetAttr); } @@ -103,7 +91,7 @@ struct TileMxNxK { int64_t K = 1; }; -MaterializeEncodingInfo +IREE::Codegen::MaterializeEncodingInfo getEncodingInfoForMatmul(IREE::Encoding::EncodingAttr encoding, int64_t rank, TileMxNxK tileMxNxK); @@ -142,10 +130,6 @@ void populateShapeIndependentMaterializeEncodingPatterns( // result of a matvec. bool isNarrowNResult(IREE::Encoding::EncodingAttr encoding); -/// Concatenates the vectors. -SmallVector -getExpandedTileShape(const TileSwizzle::ExpandShapeType &expandShape); - } // namespace mlir::iree_compiler #endif // IREE_COMPILER_SRC_IREE_COMPILER_CODEGEN_COMMON_ENCODINGUTILS_H_ diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel index 12f68a3ff8c0..63a63d311e6c 100644 --- a/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel +++ b/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel @@ -95,6 +95,7 @@ iree_compiler_cc_library( "//compiler/src/iree/compiler/Codegen/Common", "//compiler/src/iree/compiler/Codegen/Common:VectorLayoutAnalysis", "//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/Dialect/GPU/Transforms:GPUTransforms", "//compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR:IREEVectorExtDialect", diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt index 4de9d3f5fcf0..2c1acf781c4a 100644 --- a/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt +++ b/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt @@ -126,6 +126,7 @@ iree_cc_library( iree::compiler::Codegen::Common iree::compiler::Codegen::Common::VectorLayoutAnalysis iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect + iree::compiler::Codegen::Dialect::Codegen::Utils iree::compiler::Codegen::Dialect::GPU::IR::IREEGPUDialect iree::compiler::Codegen::Dialect::GPU::Transforms::GPUTransforms iree::compiler::Codegen::Dialect::VectorExt::IR::IREEVectorExtDialect diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUMaterializeEncoding.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUMaterializeEncoding.cpp index 88f6b30f6033..e32760b44215 100644 --- a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUMaterializeEncoding.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUMaterializeEncoding.cpp @@ -8,6 +8,7 @@ #include "iree/compiler/Codegen/Common/EncodingUtils.h" #include "iree/compiler/Codegen/Common/GPU/Passes.h" #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h" +#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h" @@ -41,6 +42,9 @@ namespace mlir::iree_compiler { #define GEN_PASS_DEF_GPUMATERIALIZEHOSTENCODINGPASS #include "iree/compiler/Codegen/Common/GPU/Passes.h.inc" +using IREE::Codegen::MaterializeEncodingInfo; +using IREE::Codegen::TileSwizzle; + static IREE::GPU::MMAAttr chooseIntrinsicMMAAttr(TypeRange eTypes, IREE::GPU::TargetWgpAttr wgp) { IREE::GPU::MMAAttr candidateMma; diff --git a/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoNop.cpp b/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoNop.cpp index 6a55b71be15c..32eb822c189d 100644 --- a/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoNop.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoNop.cpp @@ -22,6 +22,7 @@ namespace mlir::iree_compiler { #include "iree/compiler/Codegen/Common/Passes.h.inc" using namespace IREE::Encoding; +using IREE::Codegen::MaterializeEncodingInfo; namespace { struct MaterializeEncodingIntoNopPass final diff --git a/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoPackUnPack.cpp b/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoPackUnPack.cpp index d68273f28729..57a990b78bfc 100644 --- a/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoPackUnPack.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/MaterializeEncodingIntoPackUnPack.cpp @@ -28,6 +28,8 @@ namespace mlir::iree_compiler { +using IREE::Codegen::MaterializeEncodingInfo; + //===---------------------------------------------------------------------===// // Utility methods //===---------------------------------------------------------------------===// @@ -48,7 +50,7 @@ getSwizzledShape(ArrayRef packedShape, SmallVector newShape(packedShape.take_front(srcRank)); SmallVector expandedTileShape = - getExpandedTileShape(encodingInfo.swizzle->expandShape); + IREE::Codegen::getExpandedTileShape(encodingInfo.swizzle->expandShape); MLIRContext *ctx = packedShape[0].getContext(); Builder b(ctx); for (int64_t d : expandedTileShape) { diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/BUILD.bazel new file mode 100644 index 000000000000..fc03e8d88f77 --- /dev/null +++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/BUILD.bazel @@ -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 + +load("//build_tools/bazel:build_defs.oss.bzl", "iree_compiler_cc_library") + +package( + default_visibility = ["//visibility:public"], + features = ["layering_check"], + licenses = ["notice"], # Apache 2.0 +) + +iree_compiler_cc_library( + name = "Utils", + srcs = [ + "Utils.cpp", + ], + hdrs = [ + "Utils.h", + ], + deps = [ + "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:IR", + ], +) diff --git a/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/CMakeLists.txt new file mode 100644 index 000000000000..cbfa9245cc88 --- /dev/null +++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/CMakeLists.txt @@ -0,0 +1,27 @@ +################################################################################ +# Autogenerated by build_tools/bazel_to_cmake/bazel_to_cmake.py from # +# compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/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 + Utils + HDRS + "Utils.h" + SRCS + "Utils.cpp" + DEPS + LLVMSupport + MLIRIR + iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect + PUBLIC +) + +### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ### diff --git a/compiler/src/iree/compiler/Codegen/Common/TileSwizzle.cpp b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.cpp similarity index 63% rename from compiler/src/iree/compiler/Codegen/Common/TileSwizzle.cpp rename to compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.cpp index 289952336c76..68a9630f0267 100644 --- a/compiler/src/iree/compiler/Codegen/Common/TileSwizzle.cpp +++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.cpp @@ -4,10 +4,14 @@ // See https://llvm.org/LICENSE.txt for license information. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -#include "iree/compiler/Codegen/Common/TileSwizzle.h" +#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h" #include "llvm/ADT/STLExtras.h" -namespace mlir::iree_compiler { +namespace mlir::iree_compiler::IREE::Codegen { + +//===----------------------------------------------------------------------===// +// Layout Structs. +//===----------------------------------------------------------------------===// llvm::raw_ostream &operator<<(llvm::raw_ostream &os, TileSwizzle::Dim::Kind kind) { @@ -47,4 +51,19 @@ llvm::raw_ostream &operator<<(llvm::raw_ostream &os, return os; } -} // namespace mlir::iree_compiler +//===----------------------------------------------------------------------===// +// Layout Utilities. +//===----------------------------------------------------------------------===// + +SmallVector +getExpandedTileShape(const TileSwizzle::ExpandShapeType &expandShape) { + SmallVector result; + for (auto e : expandShape) { + for (auto d : e) { + result.push_back(d.size); + } + } + return result; +} + +} // namespace mlir::iree_compiler::IREE::Codegen diff --git a/compiler/src/iree/compiler/Codegen/Common/TileSwizzle.h b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h similarity index 72% rename from compiler/src/iree/compiler/Codegen/Common/TileSwizzle.h rename to compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h index 82eff5965521..04091a79b1a1 100644 --- a/compiler/src/iree/compiler/Codegen/Common/TileSwizzle.h +++ b/compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h @@ -4,14 +4,21 @@ // 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_COMMON_TILESWIZZLE_H_ -#define IREE_COMPILER_SRC_IREE_COMPILER_CODEGEN_COMMON_TILESWIZZLE_H_ +#ifndef IREE_COMPILER_CODEGEN_DIALECT_CODEGEN_UTILS_H_ +#define IREE_COMPILER_CODEGEN_DIALECT_CODEGEN_UTILS_H_ #include + +#include "llvm-c/TargetMachine.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/raw_ostream.h" +#include "mlir/Support/LLVM.h" + +namespace mlir::iree_compiler::IREE::Codegen { -namespace mlir::iree_compiler { +//===----------------------------------------------------------------------===// +// Layout Structs. +//===----------------------------------------------------------------------===// // Metadata for a swizzle, that is, an (expand_shape -> transposition) // pair of ops performing a change of layout within the tiles. This is used @@ -80,6 +87,27 @@ llvm::raw_ostream &operator<<(llvm::raw_ostream &os, TileSwizzle::Dim dim); llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const TileSwizzle &swizzle); -} // namespace mlir::iree_compiler +/// Container of information needed to materialize the layout transformations. +struct MaterializeEncodingInfo { + // The next 3 fields are used to create a `tensor.pack` or `tensor.unpack` op, + // changing the overall layout between row-major and tiled (where each tile is + // row-major). + SmallVector innerDimsPos; + SmallVector innerTileSizes; + SmallVector outerDimsPerm; + + // The optional swizzle, see the comment on TileSwizzle. Only used on GPU. + std::optional swizzle; +}; + +//===----------------------------------------------------------------------===// +// Layout Utilities. +//===----------------------------------------------------------------------===// + +/// Concatenates the vectors. +SmallVector +getExpandedTileShape(const TileSwizzle::ExpandShapeType &expandShape); + +} // namespace mlir::iree_compiler::IREE::Codegen -#endif // IREE_COMPILER_SRC_IREE_COMPILER_CODEGEN_COMMON_TILESWIZZLE_H_ +#endif // IREE_COMPILER_CODEGEN_DIALECT_CODEGEN_UTILS_H_ diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel index a28d668eb043..97ea44550e30 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel @@ -83,8 +83,8 @@ iree_compiler_cc_library( ":IREEGPUInterfaces", ":IREEGPUOpsGen", "//compiler/bindings/c:headers", - "//compiler/src/iree/compiler/Codegen/Common:TileSwizzle", "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect", + "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils", "//compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR:IREEVectorExtDialect", "//compiler/src/iree/compiler/Codegen/Utils:VectorOpUtils", "//compiler/src/iree/compiler/Dialect/LinalgExt/IR", diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt index 907f34bdc9e5..e0df65823007 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt @@ -63,8 +63,8 @@ iree_cc_library( MLIRTilingInterface MLIRVectorDialect MLIRVectorInterfaces - iree::compiler::Codegen::Common::TileSwizzle iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect + iree::compiler::Codegen::Dialect::Codegen::Utils iree::compiler::Codegen::Dialect::VectorExt::IR::IREEVectorExtDialect iree::compiler::Codegen::Utils::VectorOpUtils iree::compiler::Dialect::LinalgExt::IR diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.cpp index 9d0b717162db..f80c345d8513 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.cpp +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.cpp @@ -3,9 +3,13 @@ // 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/GPU/IR/GPUTileSwizzleUtils.h" -namespace mlir::iree_compiler { +#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h" + +namespace mlir::iree_compiler::IREE::GPU { +using ::mlir::iree_compiler::IREE::Codegen::TileSwizzle; using Kind = TileSwizzle::Dim::Kind; // Returns the index of the first destination dimension corresponding to the @@ -193,4 +197,4 @@ TileSwizzle getSwizzle(IREE::GPU::DataTiledMMAAttr mma, return swizzle; } -} // namespace mlir::iree_compiler +} // namespace mlir::iree_compiler::IREE::GPU diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.h b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.h index 413f80176c44..c1b9fcc44ca9 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.h +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.h @@ -1,4 +1,4 @@ -// Copyright 2023 The IREE Authors +// Copyright 2024 The IREE Authors // // Licensed under the Apache License v2.0 with LLVM Exceptions. // See https://llvm.org/LICENSE.txt for license information. @@ -7,22 +7,22 @@ #ifndef IREE_COMPILER_CODEGEN_DIALECT_GPU_IR_GPUTILESWIZZLEUTILS_H_ #define IREE_COMPILER_CODEGEN_DIALECT_GPU_IR_GPUTILESWIZZLEUTILS_H_ -#include "iree/compiler/Codegen/Common/TileSwizzle.h" +#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUEnums.h" -namespace mlir::iree_compiler { +namespace mlir::iree_compiler::IREE::GPU { // Returns the TileSwizzle bringing a tile from row-major layout into the tiled // layout consumed by the given `intrinsic` and `fragment`. -TileSwizzle getIntrinsicSwizzle(IREE::GPU::MMAIntrinsic intrinsic, - IREE::GPU::MMAFragment fragment); +Codegen::TileSwizzle getIntrinsicSwizzle(IREE::GPU::MMAIntrinsic intrinsic, + IREE::GPU::MMAFragment fragment); // Returns the swizzle for the full data-tiled-mma tile, including all the // relevant unrolling and expansion factors. -TileSwizzle getSwizzle(IREE::GPU::DataTiledMMAAttr mma, - IREE::GPU::MMAFragment fragment); +Codegen::TileSwizzle getSwizzle(IREE::GPU::DataTiledMMAAttr mma, + IREE::GPU::MMAFragment fragment); -} // namespace mlir::iree_compiler +} // namespace mlir::iree_compiler::IREE::GPU #endif // IREE_COMPILER_CODEGEN_DIALECT_GPU_IR_GPUTILESWIZZLEUTILS_H_ diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp index 93370d02b481..2fae167aa987 100644 --- a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp +++ b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp @@ -46,6 +46,8 @@ namespace mlir::iree_compiler::IREE::GPU { +using ::mlir::iree_compiler::IREE::Codegen::TileSwizzle; + //===----------------------------------------------------------------------===// // MMA intrinsics semantics: shapes, layouts, operand element types. //===----------------------------------------------------------------------===// diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/BUILD.bazel b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/BUILD.bazel index 17c96e112389..3227ac21ea14 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/BUILD.bazel +++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/BUILD.bazel @@ -60,6 +60,7 @@ iree_compiler_cc_library( "//compiler/src/iree/compiler/Codegen/Common", "//compiler/src/iree/compiler/Codegen/Common/GPU:CommonGPUPasses", "//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/LLVMGPU/Utils", "//compiler/src/iree/compiler/Codegen/Utils", diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/CMakeLists.txt index c9e2ef557ce6..f027fbe74961 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/CMakeLists.txt +++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/CMakeLists.txt @@ -65,6 +65,7 @@ iree_cc_library( iree::compiler::Codegen::Common iree::compiler::Codegen::Common::GPU::CommonGPUPasses iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect + iree::compiler::Codegen::Dialect::Codegen::Utils iree::compiler::Codegen::Dialect::GPU::IR::IREEGPUDialect iree::compiler::Codegen::LLVMGPU::Utils iree::compiler::Codegen::Utils diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.cpp index b0022c452ab1..7dc2e4093d58 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.cpp +++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/TransformExtensions/LLVMGPUExtensions.cpp @@ -11,6 +11,7 @@ #include "iree/compiler/Codegen/Common/GPU/GPUVectorDistribution.h" #include "iree/compiler/Codegen/Common/GPU/Passes.h" #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h" +#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUInterfaces.h" #include "iree/compiler/Codegen/LLVMGPU/Utils/LLVMGPUUtils.h" @@ -55,6 +56,7 @@ using llvm::dbgs; using namespace mlir; using namespace mlir::iree_compiler; using namespace mlir::iree_compiler::IREE; +using mlir::iree_compiler::IREE::Codegen::TileSwizzle; iree_compiler::IREE::transform_dialect::LLVMGPUExtensions::LLVMGPUExtensions() { // CreateAsyncGroupsOp depends on the following two dialects.