move GeneralizeNamedOp pass

compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel

@@ -51,6 +51,7 @@ iree_compiler_cc_library(
         "GPUCheckResourceUsage.cpp",
         "GPUDistribute.cpp",
         "GPUDistributeSharedMemoryCopy.cpp",
+        "GPUGeneralizeNamedOps.cpp",
         "GPUMultiBuffering.cpp",
         "GPUPatterns.cpp",
         "GPUPipelining.cpp",

compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt

@@ -49,6 +49,7 @@ iree_cc_library(
     "GPUCheckResourceUsage.cpp"
     "GPUDistribute.cpp"
     "GPUDistributeSharedMemoryCopy.cpp"
+    "GPUGeneralizeNamedOps.cpp"
     "GPUMultiBuffering.cpp"
     "GPUPatterns.cpp"
     "GPUPipelining.cpp"

compiler/src/iree/compiler/Codegen/Common/GPU/GPUGeneralizeNamedOps.cpp

@@ -0,0 +1,57 @@
+// 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
+
+//===- GPUGeneralizeNamedOps.cpp - Pass to generalize named linalg ops --===//
+//
+// The pass is to generalize named linalg ops that are better as linalg.generic
+// ops in IREE
+//
+//===----------------------------------------------------------------------===//
+
+#include "iree/compiler/Codegen/Common/GPU/PassDetail.h"
+#include "iree/compiler/Codegen/Common/GPU/Passes.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+namespace {
+struct GPUGeneralizeNamedOpsPass
+    : public GPUGeneralizeNamedOpsBase<GPUGeneralizeNamedOpsPass> {
+
+  void runOnOperation() override;
+};
+} // namespace
+
+void GPUGeneralizeNamedOpsPass::runOnOperation() {
+  auto funcOp = getOperation();
+  SmallVector<linalg::LinalgOp> namedOpCandidates;
+  funcOp.walk([&](linalg::LinalgOp linalgOp) {
+    if (isa<linalg::BatchMatmulTransposeBOp, linalg::MatmulTransposeBOp,
+            linalg::VecmatOp, linalg::MatvecOp>(linalgOp))
+      namedOpCandidates.push_back(linalgOp);
+  });
+
+  IRRewriter rewriter(&getContext());
+  for (auto linalgOp : namedOpCandidates) {
+    rewriter.setInsertionPoint(linalgOp);
+    FailureOr<linalg::GenericOp> generalizedOp =
+        linalg::generalizeNamedOp(rewriter, linalgOp);
+    if (failed(generalizedOp)) {
+      linalgOp->emitOpError("failed to generalize operation");
+      return signalPassFailure();
+    }
+  }
+}
+
+std::unique_ptr<OperationPass<func::FuncOp>> createGPUGeneralizeNamedOpsPass() {
+  return std::make_unique<GPUGeneralizeNamedOpsPass>();
+}
+
+} // namespace iree_compiler
+} // namespace mlir

compiler/src/iree/compiler/Codegen/Common/GPU/Passes.h

@@ -124,6 +124,9 @@ createWorkgroupSpecializationPass();
 std::unique_ptr<OperationPass<func::FuncOp>>
 createWorkGroupSwizzle(unsigned swizzleLogTile = 0);
 
+// This pass generalizes named Linalg ops that are better off as generics.
+std::unique_ptr<OperationPass<func::FuncOp>> createGPUGeneralizeNamedOpsPass();
+
 /// Register Common GPU passes.
 void registerCodegenCommonGPUPasses();
 

compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td

@@ -31,6 +31,12 @@ def GPUDistributeSharedMemoryCopy :
   let constructor = "mlir::iree_compiler::createGPUDistributeSharedMemoryCopy()";
 }
 
+def GPUGeneralizeNamedOps :
+    Pass<"iree-codegen-gpu-generalize-named-ops", "func::FuncOp"> {
+  let summary = "Convert named Linalg ops to linalg.generic ops";
+  let constructor = "mlir::iree_compiler::createGPUGeneralizeNamedOpsPass()";
+}
+
 def GPUMultiBuffering :
     Pass<"iree-codegen-gpu-multi-buffering", "func::FuncOp"> {
   let summary = "Pass to do multi buffering.";

compiler/src/iree/compiler/Codegen/SPIRV/test/generalize_named_ops.mlir → compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_generalize_named_ops.mlir

@@ -1,4 +1,4 @@
-// RUN: iree-opt --split-input-file --iree-spirv-generalize-named-ops %s | FileCheck %s
+// RUN: iree-opt --split-input-file --iree-codegen-gpu-generalize-named-ops %s | FileCheck %s
 
 module {
   func.func @transpose_matmul(%arg0: tensor<1x4096xf32>, %arg1: tensor<32000x4096xf32>) -> tensor<1x32000xf32> {

compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp

@@ -598,6 +598,8 @@ void addGPUTransformDialectPasses(OpPassManager &passManager) {
 
 void buildLLVMGPUCodegenConfigurationPassPipeline(OpPassManager &pm) {
   addCommonTargetExecutablePreprocessingPasses(pm);
+  auto &nestedModulePM = pm.nest<ModuleOp>();
+  nestedModulePM.addNestedPass<func::FuncOp>(createGPUGeneralizeNamedOpsPass());
   pm.addPass(createLLVMGPUSelectLoweringStrategyPass());
 }
 

compiler/src/iree/compiler/Codegen/SPIRV/BUILD.bazel

@@ -62,7 +62,6 @@ iree_compiler_cc_library(
         "SPIRVEmulateI64.cpp",
         "SPIRVEraseStorageBufferStaticShape.cpp",
         "SPIRVFinalVectorLowering.cpp",
-        "SPIRVGeneralizeNamedOps.cpp",
         "SPIRVInitialVectorLowering.cpp",
         "SPIRVLowerExecutableTargetPass.cpp",
         "SPIRVMapMemRefStorageClass.cpp",

compiler/src/iree/compiler/Codegen/SPIRV/CMakeLists.txt

@@ -61,7 +61,6 @@ iree_cc_library(
     "SPIRVEmulateI64.cpp"
     "SPIRVEraseStorageBufferStaticShape.cpp"
     "SPIRVFinalVectorLowering.cpp"
-    "SPIRVGeneralizeNamedOps.cpp"
     "SPIRVInitialVectorLowering.cpp"
     "SPIRVLowerExecutableTargetPass.cpp"
     "SPIRVMapMemRefStorageClass.cpp"

compiler/src/iree/compiler/Codegen/SPIRV/Passes.cpp

@@ -664,8 +664,7 @@ void addSPIRVTransformDialectPassPipeline(OpPassManager &pm) {
 void buildSPIRVCodegenConfigurationPassPipeline(OpPassManager &pm) {
   addCommonTargetExecutablePreprocessingPasses(pm);
   auto &nestedModulePM = pm.nest<ModuleOp>();
-  nestedModulePM.addNestedPass<func::FuncOp>(
-      createSPIRVGeneralizeNamedOpsPass());
+  nestedModulePM.addNestedPass<func::FuncOp>(createGPUGeneralizeNamedOpsPass());
   pm.addPass(createSPIRVSelectLoweringStrategyPass());
 }
 

compiler/src/iree/compiler/Codegen/SPIRV/Passes.h

@@ -102,10 +102,6 @@ createSPIRVEraseStorageBufferStaticShapePass();
 std::unique_ptr<OperationPass<func::FuncOp>>
 createSPIRVFoldProcessorIDUsesPass();
 
-// This pass generalizes named Linalg ops that are better off as generics.
-std::unique_ptr<OperationPass<func::FuncOp>>
-createSPIRVGeneralizeNamedOpsPass();
-
 /// Pass to set the lowering strategy for the target variant.
 std::unique_ptr<OperationPass<IREE::HAL::ExecutableVariantOp>>
 createSPIRVSelectLoweringStrategyPass();

compiler/src/iree/compiler/Codegen/SPIRV/Passes.td

@@ -58,12 +58,6 @@ def SPIRVEraseStorageBufferStaticShape :
   let constructor = "mlir::iree_compiler::createSPIRVEraseStorageBufferStaticShapePass()";
 }
 
-def SPIRVGeneralizeNamedOps :
-    Pass<"iree-spirv-generalize-named-ops", "func::FuncOp"> {
-  let summary = "Convert named Linalg ops to linalg.generic ops";
-  let constructor = "mlir::iree_compiler::createSPIRVGeneralizeNamedOpsPass()";
-}
-
 def SPIRVLowerExecutableTarget :
     Pass<"iree-spirv-lower-executable-target-pass",
          "mlir::iree_compiler::IREE::HAL::ExecutableVariantOp"> {