From 55b998a5f09c9365cdaed4edc55a06b01f5a577c Mon Sep 17 00:00:00 2001 From: Bangtian Liu Date: Mon, 11 Nov 2024 15:42:20 -0500 Subject: [PATCH] [tuner] add an iree-opt pass to strip configuration from executable sources (#19069) This PR aims to address the first task in https://github.com/nod-ai/SHARK-Platform/issues/453: adding an iree-opt pass that removes configuration from executable sources. The corresponding test is also included to ensure its correct functionality. --------- Signed-off-by: Bangtian Liu --- .../iree/compiler/Codegen/Common/BUILD.bazel | 1 + .../compiler/Codegen/Common/CMakeLists.txt | 1 + .../iree/compiler/Codegen/Common/Passes.td | 5 ++ .../Common/StripCompilationInfoPass.cpp | 71 +++++++++++++++++++ .../compiler/Codegen/Common/test/BUILD.bazel | 1 + .../Codegen/Common/test/CMakeLists.txt | 1 + .../Common/test/strip_compilation_info.mlir | 62 ++++++++++++++++ 7 files changed, 142 insertions(+) create mode 100644 compiler/src/iree/compiler/Codegen/Common/StripCompilationInfoPass.cpp create mode 100644 compiler/src/iree/compiler/Codegen/Common/test/strip_compilation_info.mlir diff --git a/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel index 4a6afa18252c..cc43b0918bad 100644 --- a/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel +++ b/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel @@ -139,6 +139,7 @@ iree_compiler_cc_library( "RemoveSingleIterationLoop.cpp", "ReplaceSlowMinMaxOps.cpp", "SplitFullPartialTransferPass.cpp", + "StripCompilationInfoPass.cpp", "TensorDynamicDimAnalysis.cpp", "TensorToVectorVectorizePad.cpp", "TestExecutablePreprocessing.cpp", diff --git a/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt index eb56df21c8e0..c87a67d3358b 100644 --- a/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt +++ b/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt @@ -131,6 +131,7 @@ iree_cc_library( "RemoveSingleIterationLoop.cpp" "ReplaceSlowMinMaxOps.cpp" "SplitFullPartialTransferPass.cpp" + "StripCompilationInfoPass.cpp" "TensorDynamicDimAnalysis.cpp" "TensorToVectorVectorizePad.cpp" "TestExecutablePreprocessing.cpp" diff --git a/compiler/src/iree/compiler/Codegen/Common/Passes.td b/compiler/src/iree/compiler/Codegen/Common/Passes.td index b86faabdfd05..a649f3703ae8 100644 --- a/compiler/src/iree/compiler/Codegen/Common/Passes.td +++ b/compiler/src/iree/compiler/Codegen/Common/Passes.td @@ -490,6 +490,11 @@ def RemoveSingleIterationLoopPass : let summary = "Remove distributed loop with single iteration."; } +def StripCompilationInfoPass : + Pass<"iree-codegen-strip-compilation-info", "">{ + let summary = "Remove all the the lowering configuration and translation info attributes."; +} + // TODO: Replace with upstream: https://github.com/iree-org/iree/issues/18759 def IREELoopInvariantCodeMotionPass : Pass<"iree-loop-invariant-code-motion", ""> { diff --git a/compiler/src/iree/compiler/Codegen/Common/StripCompilationInfoPass.cpp b/compiler/src/iree/compiler/Codegen/Common/StripCompilationInfoPass.cpp new file mode 100644 index 000000000000..fe06412bd755 --- /dev/null +++ b/compiler/src/iree/compiler/Codegen/Common/StripCompilationInfoPass.cpp @@ -0,0 +1,71 @@ +// 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 + +#include "iree/compiler/Codegen/Common/Passes.h" +#include "mlir/Transforms/WalkPatternRewriteDriver.h" + +#define DEBUG_TYPE "iree-codegen-strip-compilation-info" + +namespace mlir::iree_compiler { + +#define GEN_PASS_DEF_STRIPCOMPILATIONINFOPASS +#include "iree/compiler/Codegen/Common/Passes.h.inc" + +namespace { + +struct StripFuncOpTranslationInfo final + : OpInterfaceRewritePattern { + using OpInterfaceRewritePattern::OpInterfaceRewritePattern; + LogicalResult matchAndRewrite(mlir::FunctionOpInterface funcOp, + PatternRewriter &rewriter) const final { + if (!getTranslationInfo(funcOp)) + return failure(); + + rewriter.modifyOpInPlace(funcOp, [&]() { + // If the function has translation info, erase it. + eraseTranslationInfo(funcOp); + }); + + return success(); + } +}; + +struct StripLinalgOpCompilationInfo final + : OpInterfaceRewritePattern { + using OpInterfaceRewritePattern::OpInterfaceRewritePattern; + LogicalResult matchAndRewrite(linalg::LinalgOp linalgOp, + PatternRewriter &rewriter) const final { + if (!getCompilationInfo(linalgOp) && !getLoweringConfig(linalgOp)) + return failure(); + + rewriter.modifyOpInPlace(linalgOp, [&]() { + if (getCompilationInfo(linalgOp)) { + // Erase the compilation info configuration if it exists. + eraseCompilationInfo(linalgOp); + } + if (getLoweringConfig(linalgOp)) { + // Erase the lowering configuration from root operation if it + // exists. + eraseLoweringConfig(linalgOp); + } + }); + + return success(); + } +}; + +struct StripCompilationInfoPass final + : impl::StripCompilationInfoPassBase { + void runOnOperation() override { + MLIRContext *ctx = &getContext(); + RewritePatternSet patterns(ctx); + patterns.add(ctx); + patterns.add(ctx); + walkAndApplyPatterns(getOperation(), std::move(patterns)); + } +}; +} // namespace +} // namespace mlir::iree_compiler diff --git a/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel index 403975b20174..0c3b4349238a 100644 --- a/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel +++ b/compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel @@ -67,6 +67,7 @@ iree_lit_test_suite( "remove_trivial_loops.mlir", "repeated_matcher_use.mlir", "replace_slow_min_max_ops.mlir", + "strip_compilation_info.mlir", "test_partitionable_loops_interface.mlir", "tile_and_distribute_to_workgroups_func_scope.mlir", "tile_and_distribute_to_workgroups.mlir", diff --git a/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt index 2156340f2a87..2bfcf3fdfd41 100644 --- a/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt +++ b/compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt @@ -63,6 +63,7 @@ iree_lit_test_suite( "remove_trivial_loops.mlir" "repeated_matcher_use.mlir" "replace_slow_min_max_ops.mlir" + "strip_compilation_info.mlir" "test_partitionable_loops_interface.mlir" "tile_and_distribute_to_workgroups.mlir" "tile_and_distribute_to_workgroups_func_scope.mlir" diff --git a/compiler/src/iree/compiler/Codegen/Common/test/strip_compilation_info.mlir b/compiler/src/iree/compiler/Codegen/Common/test/strip_compilation_info.mlir new file mode 100644 index 000000000000..a08f187a4405 --- /dev/null +++ b/compiler/src/iree/compiler/Codegen/Common/test/strip_compilation_info.mlir @@ -0,0 +1,62 @@ +// RUN: iree-opt --split-input-file --iree-codegen-strip-compilation-info %s | FileCheck %s + +#translation_info = #iree_codegen.translation_info +func.func @main() attributes {translation_info = #translation_info} { + return +} + +// CHECK-LABEL: func.func @main +// CHECK-NOT: iree_codegen.translation_info +// CHECK-NOT: LLVMGPUVectorDistribute + +// ----- + +#pipeline_layout = #hal.pipeline.layout +]> +hal.executable private @strip_main { + hal.executable.variant public @strip_main target(#hal.executable.target<"", "", {}>) { + hal.executable.export public @entry_point layout(#pipeline_layout) + builtin.module { + func.func @fn1() attributes {translation_info = #iree_codegen.translation_info} { + return + } + func.func @fn2() attributes {translation_info = #iree_codegen.translation_info} { + return + } + } + } +} + +// CHECK-LABEL: hal.executable private @strip_main +// CHECK: @fn1 +// CHECK-NOT: translation_info = +// CHECK: @fn2 +// CHECK-NOT: translation_info = +// CHECK: return + +// ----- + +#config = #iree_codegen.lowering_config +#translation = #iree_codegen.translation_info +#compilation = #iree_codegen.compilation_info +func.func @matmul_128x1024x256(%lhs : tensor<128x256xf32>, %rhs: tensor<256x1024xf32>, %init: tensor<128x1024xf32>) -> tensor<128x1024xf32> { + %result = linalg.matmul {compilation_info = #compilation} ins(%lhs, %rhs : tensor<128x256xf32>, tensor<256x1024xf32>) outs(%init : tensor<128x1024xf32>) -> tensor<128x1024xf32> + return %result : tensor<128x1024xf32> +} + +// CHECK-LABEL: func.func @matmul_128x1024x256 +// CHECK-NOT: iree_codegen.translation_info +// CHECK-NOT: iree_codegen.lowering_config +// CHECK-NOT: iree_codegen.compilation_info + +// ----- + +#config = #iree_codegen.lowering_config +func.func @matmul_128x1024x256_1(%lhs : tensor<128x256xf32>, %rhs: tensor<256x1024xf32>, %init: tensor<128x1024xf32>) -> tensor<128x1024xf32> { + %result = linalg.matmul {lowering_config = #iree_codegen.lowering_config} ins(%lhs, %rhs : tensor<128x256xf32>, tensor<256x1024xf32>) outs(%init : tensor<128x1024xf32>) -> tensor<128x1024xf32> + return %result : tensor<128x1024xf32> +} + +// CHECK-LABEL: func.func @matmul_128x1024x256_1 +// CHECK-NOT: iree_codegen.lowering_config