[Codegen][NFC] Move encoding related containers to Dialect/Codegen/Ut…

…ils (iree-org#19234)

The revision moves MaterializeEncodingInfo struct and TileSwizzle struct
to `compiler/Codegen/Dialect/Codegen/Utils/Utils.[h|cpp]`. It is a
preparation for iree-org#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:
iree-org@740e301

---------

Signed-off-by: hanhanW <hanhan0912@gmail.com>

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

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

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",

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

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"

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<int64_t> newShape(
         packedType.getShape().drop_back(encodingInfo.innerTileSizes.size()));
     SmallVector<int64_t> 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<int64_t>
-getExpandedTileShape(const TileSwizzle::ExpandShapeType &expandShape) {
-  SmallVector<int64_t> result;
-  for (auto e : expandShape) {
-    for (auto d : e) {
-      result.push_back(d.size);
-    }
-  }
-  return result;
-}
-
 } // namespace mlir::iree_compiler

compiler/src/iree/compiler/Codegen/Common/EncodingUtils.h

@@ -7,29 +7,17 @@
 #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"
 #include "mlir/Transforms/DialectConversion.h"
 
 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<int64_t> innerDimsPos;
-  SmallVector<int64_t> innerTileSizes;
-  SmallVector<int64_t> outerDimsPerm;
-
-  // The optional swizzle, see the comment on TileSwizzle. Only used on GPU.
-  std::optional<TileSwizzle> swizzle;
-};
-
-using MaterializeEncodingFn = std::function<FailureOr<MaterializeEncodingInfo>(
-    RankedTensorType, IREE::HAL::ExecutableTargetAttr targetAttr)>;
+using MaterializeEncodingFn =
+    std::function<FailureOr<IREE::Codegen::MaterializeEncodingInfo>(
+        RankedTensorType, IREE::HAL::ExecutableTargetAttr targetAttr)>;
 
 struct MaterializeEncodingValueInfo {
   SmallVector<Value> innerTileSizes;
@@ -56,7 +44,7 @@ class MaterializeEncodingTypeConverter : public TypeConverter {
 
   IREE::HAL::ExecutableTargetAttr getTargetAttr() const { return targetAttr; }
 
-  FailureOr<MaterializeEncodingInfo>
+  FailureOr<IREE::Codegen::MaterializeEncodingInfo>
   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<int64_t>
-getExpandedTileShape(const TileSwizzle::ExpandShapeType &expandShape);
-
 } // namespace mlir::iree_compiler
 
 #endif // IREE_COMPILER_SRC_IREE_COMPILER_CODEGEN_COMMON_ENCODINGUTILS_H_

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",

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

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;

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

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<OpFoldResult> packedShape,
 
   SmallVector<OpFoldResult> newShape(packedShape.take_front(srcRank));
   SmallVector<int64_t> expandedTileShape =
-      getExpandedTileShape(encodingInfo.swizzle->expandShape);
+      IREE::Codegen::getExpandedTileShape(encodingInfo.swizzle->expandShape);
   MLIRContext *ctx = packedShape[0].getContext();
   Builder b(ctx);
   for (int64_t d : expandedTileShape) {

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",
+    ],
+)

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 ###

compiler/src/iree/compiler/Codegen/Common/TileSwizzle.cpp → 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<int64_t>
+getExpandedTileShape(const TileSwizzle::ExpandShapeType &expandShape) {
+  SmallVector<int64_t> result;
+  for (auto e : expandShape) {
+    for (auto d : e) {
+      result.push_back(d.size);
+    }
+  }
+  return result;
+}
+
+} // namespace mlir::iree_compiler::IREE::Codegen

compiler/src/iree/compiler/Codegen/Common/TileSwizzle.h → 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 <cstdint>
+
+#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<int64_t> innerDimsPos;
+  SmallVector<int64_t> innerTileSizes;
+  SmallVector<int64_t> outerDimsPerm;
+
+  // The optional swizzle, see the comment on TileSwizzle. Only used on GPU.
+  std::optional<TileSwizzle> swizzle;
+};
+
+//===----------------------------------------------------------------------===//
+// Layout Utilities.
+//===----------------------------------------------------------------------===//
+
+/// Concatenates the vectors.
+SmallVector<int64_t>
+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_