Lower data-tiled multi_mma to intrinsics. (iree-org#18547)

This lowers `multi_mma` with `DataTiledMMAAttr` down to `amdgpu` intrinsics. --------- Signed-off-by: Benoit Jacob <jacob.benoit.1@gmail.com>
nod-ai · Sep 25, 2024 · 3773a48 · 3773a48
1 parent 129ad45
commit 3773a48
Show file tree

Hide file tree

Showing 13 changed files with 339 additions and 25 deletions.
diff --git a/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel
@@ -143,7 +143,6 @@ iree_compiler_cc_library(
         "TileDispatchUsingForall.cpp",
         "TileDispatchUsingInterface.cpp",
         "TileSizeSelection.cpp",
-        "TileSwizzle.cpp",
         "TypePropagationPass.cpp",
         "UserConfig.cpp",
         "VectorizeMemrefCopy.cpp",
@@ -155,14 +154,14 @@ iree_compiler_cc_library(
         "PassUtils.h",
         "Passes.h",
         "TileSizeSelection.h",
-        "TileSwizzle.h",
         "Transforms.h",
         "UserConfig.h",
     ],
     deps = [
         ":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/VectorExt/IR:IREEVectorExtDialect",
         "//compiler/src/iree/compiler/Codegen/Interfaces:BufferizationInterfaces",
@@ -229,6 +228,13 @@ 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
@@ -73,7 +73,6 @@ iree_cc_library(
     "PassUtils.h"
     "Passes.h"
     "TileSizeSelection.h"
-    "TileSwizzle.h"
     "Transforms.h"
     "UserConfig.h"
   SRCS
@@ -135,7 +134,6 @@ iree_cc_library(
     "TileDispatchUsingForall.cpp"
     "TileDispatchUsingInterface.cpp"
     "TileSizeSelection.cpp"
-    "TileSwizzle.cpp"
     "TypePropagationPass.cpp"
     "UserConfig.cpp"
     "VectorizeMemrefCopy.cpp"
@@ -190,6 +188,7 @@ 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::VectorExt::IR::IREEVectorExtDialect
     iree::compiler::Codegen::Interfaces::BufferizationInterfaces
@@ -207,6 +206,18 @@ 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/GPU/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel
@@ -76,7 +76,6 @@ iree_compiler_cc_library(
         "GPUTensorTileToSerialLoops.cpp",
         "GPUTile.cpp",
         "GPUTileReduction.cpp",
-        "GPUTileSwizzleUtils.cpp",
         "GPUVectorAlloc.cpp",
         "GPUVectorDistribution.cpp",
         "GPUVerifyDistribution.cpp",
@@ -86,7 +85,6 @@ iree_compiler_cc_library(
     ],
     hdrs = [
         "GPUPatterns.h",
-        "GPUTileSwizzleUtils.h",
         "GPUVectorDistribution.h",
         "Passes.h",
     ],

diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt
@@ -44,7 +44,6 @@ iree_cc_library(
     CommonGPUPasses
   HDRS
     "GPUPatterns.h"
-    "GPUTileSwizzleUtils.h"
     "GPUVectorDistribution.h"
     "Passes.h"
   SRCS
@@ -75,7 +74,6 @@ iree_cc_library(
     "GPUTensorTileToSerialLoops.cpp"
     "GPUTile.cpp"
     "GPUTileReduction.cpp"
-    "GPUTileSwizzleUtils.cpp"
     "GPUVectorAlloc.cpp"
     "GPUVectorDistribution.cpp"
     "GPUVerifyDistribution.cpp"

diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUMaterializeEncoding.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUMaterializeEncoding.cpp
@@ -5,9 +5,9 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 #include "iree/compiler/Codegen/Common/EncodingUtils.h"
-#include "iree/compiler/Codegen/Common/GPU/GPUTileSwizzleUtils.h"
 #include "iree/compiler/Codegen/Common/GPU/Passes.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.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"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUEnums.h"
@@ -55,16 +55,16 @@ chooseDataTiledMMAAttr(TypeRange elementTypes, IREE::GPU::TargetAttr target) {
   Type lhs = elementTypes[0];
   Type rhs = elementTypes[1];
   Type out = elementTypes[2];
-  auto match = [=](MMAIntrinsic intrinsic, int unrollM, int unrollMToThreads,
-                   int unrollN, int unrollNToThreads,
+  auto match = [=](MMAIntrinsic intrinsic, int unrollM, int unrollMToSubgroups,
+                   int unrollN, int unrollNToSubgroups,
                    int unrollK) -> std::optional<DataTiledMMAAttr> {
     if (!hasIntrinsic(target, intrinsic)) {
       return std::nullopt;
     }
     auto candidate = DataTiledMMAAttr::get(
         ctx, MMAIntrinsicAttr::get(ctx, intrinsic), /*unroll_m=*/unrollM,
-        /*unroll_m_to_subgroups=*/unrollMToThreads, /*unroll_n=*/unrollN,
-        /*unroll_n_to_subgroups=*/unrollNToThreads, /*unroll_k=*/unrollK);
+        /*unroll_m_to_subgroups=*/unrollMToSubgroups, /*unroll_n=*/unrollN,
+        /*unroll_n_to_subgroups=*/unrollNToSubgroups, /*unroll_k=*/unrollK);
     auto [candidateLhs, candidateRhs, candidateOut] =
         candidate.getABCElementTypes();
     if (candidateLhs != lhs || candidateRhs != rhs || candidateOut != out) {

diff --git a/compiler/src/iree/compiler/Codegen/Common/TileSwizzle.cpp b/compiler/src/iree/compiler/Codegen/Common/TileSwizzle.cpp
@@ -18,6 +18,10 @@ llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
     return os << "CrossThread";
   case TileSwizzle::Dim::Kind::CrossIntrinsic:
     return os << "CrossIntrinsic";
+  default:
+    // Required by GCC.
+    assert(false);
+    return os;
   }
 }
 

diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/BUILD.bazel
@@ -48,6 +48,7 @@ iree_compiler_cc_library(
     name = "IREEGPUDialect",
     srcs = [
         "DerivedConfigUtils.cpp",
+        "GPUTileSwizzleUtils.cpp",
         "IREEGPUAttrs.cpp",
         "IREEGPUDialect.cpp",
         "IREEGPUInterfaces.cpp",
@@ -56,6 +57,7 @@ iree_compiler_cc_library(
     ],
     hdrs = [
         "DerivedConfigUtils.h",
+        "GPUTileSwizzleUtils.h",
         "IREEGPUAttrs.h",
         "IREEGPUDialect.h",
         "IREEGPUEnums.h",
@@ -80,6 +82,7 @@ iree_compiler_cc_library(
         ":IREEGPUEnums",
         ":IREEGPUInterfaces",
         ":IREEGPUOpsGen",
+        "//compiler/src/iree/compiler/Codegen/Common:TileSwizzle",
         "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
         "//compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR:IREEVectorExtDialect",
         "//compiler/src/iree/compiler/Codegen/Utils:VectorOpUtils",

diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/CMakeLists.txt
@@ -15,6 +15,7 @@ iree_cc_library(
     IREEGPUDialect
   HDRS
     "DerivedConfigUtils.h"
+    "GPUTileSwizzleUtils.h"
     "IREEGPUAttrs.h"
     "IREEGPUDialect.h"
     "IREEGPUEnums.h"
@@ -33,6 +34,7 @@ iree_cc_library(
     "IREEGPUOps.h.inc"
   SRCS
     "DerivedConfigUtils.cpp"
+    "GPUTileSwizzleUtils.cpp"
     "IREEGPUAttrs.cpp"
     "IREEGPUDialect.cpp"
     "IREEGPUInterfaces.cpp"
@@ -59,6 +61,7 @@ iree_cc_library(
     MLIRTilingInterface
     MLIRVectorDialect
     MLIRVectorInterfaces
+    iree::compiler::Codegen::Common::TileSwizzle
     iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
     iree::compiler::Codegen::Dialect::VectorExt::IR::IREEVectorExtDialect
     iree::compiler::Codegen::Utils::VectorOpUtils

diff --git a/...odegen/Common/GPU/GPUTileSwizzleUtils.cpp → ...en/Dialect/GPU/IR/GPUTileSwizzleUtils.cpp b/...odegen/Common/GPU/GPUTileSwizzleUtils.cpp → ...en/Dialect/GPU/IR/GPUTileSwizzleUtils.cpp
@@ -4,7 +4,7 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-#include "iree/compiler/Codegen/Common/GPU/GPUTileSwizzleUtils.h"
+#include "iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.h"
 
 namespace mlir::iree_compiler {
 
@@ -106,8 +106,8 @@ TileSwizzle getIntrinsicSwizzle(IREE::GPU::MMAIntrinsic intrinsic,
   }
   // The layout strides decide the initial swizzle.permutation.
   // Some WMMA intrinsics have tstrides=0 values, assert on that as that
-  // would defeat this algorithm. We'll need to solve that if and when we want
-  // to support data tiling on WMMA intrinsics.
+  // would defeat this algorithm.
+  // TODO(bjacob): Resolve that to support WMMA intrinsics.
   for (auto s : layout.tstrides) {
     (void)s;
     assert(s != 0);

diff --git a/.../Codegen/Common/GPU/GPUTileSwizzleUtils.h → ...egen/Dialect/GPU/IR/GPUTileSwizzleUtils.h b/.../Codegen/Common/GPU/GPUTileSwizzleUtils.h → ...egen/Dialect/GPU/IR/GPUTileSwizzleUtils.h
@@ -4,8 +4,8 @@
 // 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_GPU_GPUTILESWIZZLEUTILS_H_
-#define IREE_COMPILER_SRC_IREE_COMPILER_CODEGEN_COMMON_GPU_GPUTILESWIZZLEUTILS_H_
+#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/GPU/IR/IREEGPUAttrs.h"
@@ -55,4 +55,4 @@ void interleave(TileSwizzle &swizzle, int srcIndex,
 
 } // namespace mlir::iree_compiler
 
-#endif // IREE_COMPILER_SRC_IREE_COMPILER_CODEGEN_COMMON_GPU_GPUTILESWIZZLEUTILS_H_
+#endif // IREE_COMPILER_CODEGEN_DIALECT_GPU_IR_GPUTILESWIZZLEUTILS_H_