Making iree-util-strip-and-splat-constants use byte_patterns. (#14627)

Had to remove the path someone added for arith.constant as that can't take byte patterns (and it wasn't the intent of the pass). If someone wants that they can hoist into globals first before running this pass with `--iree-util-outline-constants`. A common usage of this would be to take a giant mlir file and run it through iree-opt with `--iree-util-outline-constants --iree-util-strip-and-splat-constants` to outline any arith.constant values into globals and then strip them. Note that the final VMFB will contain the same number of bytes as the original unstripped file.
iree-org · Aug 10, 2023 · ca51caa · ca51caa
1 parent 038d2ba
commit ca51caa
Show file tree

Hide file tree

Showing 26 changed files with 112 additions and 184 deletions.
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Flow/Transforms/BUILD.bazel
@@ -67,7 +67,6 @@ iree_compiler_cc_library(
         "RemoveZeroExtentTensors.cpp",
         "SetEncoding.cpp",
         "SplitReduction.cpp",
-        "StripAndSplatConstantVariables.cpp",
         "StripSignedness.cpp",
         "TensorPadToTensorInsertSlice.cpp",
         "VerifyInputLegality.cpp",

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
@@ -66,7 +66,6 @@ iree_cc_library(
     "RemoveZeroExtentTensors.cpp"
     "SetEncoding.cpp"
     "SplitReduction.cpp"
-    "StripAndSplatConstantVariables.cpp"
     "StripSignedness.cpp"
     "TensorPadToTensorInsertSlice.cpp"
     "VerifyInputLegality.cpp"

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.h b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.h
@@ -254,14 +254,6 @@ std::unique_ptr<Pass> createCollapseDimsPass();
 // Simplification and Development Tools
 //===----------------------------------------------------------------------===//
 
-// Strips constant util.globals and replaces them with splats.
-// This destructively removes data (often model weights and other parameters)
-// and is intended for use as a development tool.
-// TODO(scotttodd): pass pipeline with this and other development passes to
-//                  generate test cases / models suitable for check-in
-std::unique_ptr<OperationPass<mlir::ModuleOp>>
-createStripAndSplatConstantVariablesPass();
-
 /// Creates a pass to dump a graph for dispatches
 std::unique_ptr<Pass>
 createDumpDispatchGraphPass(raw_ostream &os = llvm::errs());

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td b/compiler/src/iree/compiler/Dialect/Flow/Transforms/Passes.td
@@ -109,7 +109,7 @@ def CloneProducersIntoDispatchRegions :
 def CollapseDimensions :
     InterfacePass<"iree-flow-collapse-dimensions", "mlir::FunctionOpInterface"> {
   let summary = "Collapse dimensions of Linalg Ops on tensor ops.";
-  let constructor = "mlir::iree_compiler::IREE::Flow::createCollapseDimensionsPass()";  
+  let constructor = "mlir::iree_compiler::IREE::Flow::createCollapseDimensionsPass()";
 }
 
 def DispatchWithTransformDialect :
@@ -317,12 +317,6 @@ def StripSignedness :
   let constructor = "mlir::iree_compiler::IREE::Flow::createStripSignednessPass()";
 }
 
-def StripAndSplatConstantVariables :
-    Pass<"iree-flow-strip-and-splat-constant-variables", "mlir::ModuleOp"> {
-  let summary = "Strips constant util.globals and replaces them with splats.";
-  let constructor = "mlir::iree_compiler::IREE::Flow::createStripAndSplatConstantVariablesPass()";
-}
-
 def VerifyInputLegality: Pass<"iree-verify-input-legality", ""> {
   let summary = "Checks the legality of the IR at the start of IREE flow transformation pipeline.";
   let constructor = "mlir::iree_compiler::IREE::Flow::createVerifyInputLegalityPass()";

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/StripAndSplatConstantVariables.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/StripAndSplatConstantVariables.cpp
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/BUILD.bazel
@@ -52,7 +52,6 @@ iree_lit_test_suite(
             "raise_special_ops.mlir",
             "remove_zero_extent_tensors.mlir",
             "set_encoding.mlir",
-            "strip_and_splat_constant_variables.mlir",
             "strip_signedness.mlir",
             "tensor_pad_to_tensor_insert_slice.mlir",
             "transform_dispatch_region_formation.mlir",

diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Flow/Transforms/test/CMakeLists.txt
@@ -50,7 +50,6 @@ iree_lit_test_suite(
     "raise_special_ops.mlir"
     "remove_zero_extent_tensors.mlir"
     "set_encoding.mlir"
-    "strip_and_splat_constant_variables.mlir"
     "strip_signedness.mlir"
     "tensor_pad_to_tensor_insert_slice.mlir"
     "transform_dispatch_region_formation.mlir"

diff --git a/...er/src/iree/compiler/Dialect/Flow/Transforms/test/strip_and_splat_constant_variables.mlir b/...er/src/iree/compiler/Dialect/Flow/Transforms/test/strip_and_splat_constant_variables.mlir
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Stream/Transforms/BUILD.bazel
@@ -27,7 +27,6 @@ iree_compiler_cc_library(
         "LayoutSlices.cpp",
         "MaterializeBuiltins.cpp",
         "MaterializeCopyOnWrite.cpp",
-        "OutlineConstants.cpp",
         "PackAllocations.cpp",
         "PackConstants.cpp",
         "PackDispatchOperands.cpp",

diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Stream/Transforms/CMakeLists.txt
@@ -29,7 +29,6 @@ iree_cc_library(
     "LayoutSlices.cpp"
     "MaterializeBuiltins.cpp"
     "MaterializeCopyOnWrite.cpp"
-    "OutlineConstants.cpp"
     "PackAllocations.cpp"
     "PackConstants.cpp"
     "PackDispatchOperands.cpp"

diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.cpp
@@ -68,7 +68,7 @@ void buildStreamTensorPassPipeline(OpPassManager &passManager,
   // Turn all constant ops into global variables and fix up the IR.
   // As many locations change and constants are deduplicated we'll end up with
   // a lot of extraneous IR (mostly global loads) and clean those up here.
-  passManager.addPass(IREE::Stream::createOutlineConstantsPass());
+  passManager.addPass(IREE::Util::createOutlineConstantsPass());
 
   // Perform cleanup after constant simplification as more canonicalizers may be
   // able to kick in.

diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.h b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.h
@@ -91,12 +91,6 @@ void buildStreamOptimizationPassPipeline(
 
 void registerStreamTransformPassPipelines();
 
-//===----------------------------------------------------------------------===//
-// Optimizations
-//===----------------------------------------------------------------------===//
-
-std::unique_ptr<OperationPass<mlir::ModuleOp>> createOutlineConstantsPass();
-
 //===----------------------------------------------------------------------===//
 // Conversion
 //===----------------------------------------------------------------------===//

diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td b/compiler/src/iree/compiler/Dialect/Stream/Transforms/Passes.td
@@ -9,18 +9,6 @@
 
 include "mlir/Pass/PassBase.td"
 
-//===----------------------------------------------------------------------===//
-// Optimizations
-//===----------------------------------------------------------------------===//
-
-def OutlineConstants :
-    Pass<"iree-stream-outline-constants", "mlir::ModuleOp"> {
-  let summary = "Outlines tensor constants into util.globals at the module level.";
-  let constructor = [{
-    mlir::iree_compiler::IREE::Stream::createOutlineConstantsPass()
-  }];
-}
-
 //===----------------------------------------------------------------------===//
 // Conversion
 //===----------------------------------------------------------------------===//

diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel
@@ -34,7 +34,6 @@ iree_lit_test_suite(
             "layout_slices.mlir",
             "materialize_builtins.mlir",
             "materialize_copy_on_write.mlir",
-            "outline_constants.mlir",
             "pack_allocations.mlir",
             "pack_constants.mlir",
             "pack_dispatch_operands.mlir",

diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt
@@ -32,7 +32,6 @@ iree_lit_test_suite(
     "layout_slices.mlir"
     "materialize_builtins.mlir"
     "materialize_copy_on_write.mlir"
-    "outline_constants.mlir"
     "pack_allocations.mlir"
     "pack_constants.mlir"
     "pack_dispatch_operands.mlir"

diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Util/Transforms/BUILD.bazel
@@ -26,11 +26,13 @@ iree_compiler_cc_library(
         "HoistIntoGlobals.cpp",
         "IPO.cpp",
         "ImportResources.cpp",
+        "OutlineConstants.cpp",
         "PassDetail.h",
         "Passes.cpp",
         "Patterns.cpp",
         "PropagateSubranges.cpp",
         "SimplifyGlobalAccesses.cpp",
+        "StripAndSplatConstants.cpp",
         "StripDebugOps.cpp",
         "TestConversion.cpp",
         "TestFloatRangeAnalysis.cpp",

diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Util/Transforms/CMakeLists.txt
@@ -29,11 +29,13 @@ iree_cc_library(
     "HoistIntoGlobals.cpp"
     "IPO.cpp"
     "ImportResources.cpp"
+    "OutlineConstants.cpp"
     "PassDetail.h"
     "Passes.cpp"
     "Patterns.cpp"
     "PropagateSubranges.cpp"
     "SimplifyGlobalAccesses.cpp"
+    "StripAndSplatConstants.cpp"
     "StripDebugOps.cpp"
     "TestConversion.cpp"
     "TestFloatRangeAnalysis.cpp"

diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/HoistIntoGlobals.cpp b/compiler/src/iree/compiler/Dialect/Util/Transforms/HoistIntoGlobals.cpp
@@ -28,8 +28,8 @@ namespace {
 // Maps an original value in the program to the symbol name of a global.
 using HoistedValueMap = llvm::DenseMap<Value, GlobalOp>;
 
-// expressions into globals. It is not expected that such a greedy algorithm
-// is great, but it is simple. Naive use of this algorithm very likely
+// Hoist expressions into globals. It is not expected that such a greedy
+// algorithm is great, but it is simple. Naive use of this algorithm very likely
 // favors programs that consume more memory at runtime than is strictly
 // necessary. Either this algorithm can be made smarter or a follow-on pass
 // can sink globals into the program where it is profitable to reduce

diff --git a/...ct/Stream/Transforms/OutlineConstants.cpp → ...lect/Util/Transforms/OutlineConstants.cpp b/...ct/Stream/Transforms/OutlineConstants.cpp → ...lect/Util/Transforms/OutlineConstants.cpp
@@ -6,12 +6,10 @@
 
 #include <utility>
 
-#include "iree/compiler/Dialect/Stream/IR/StreamDialect.h"
-#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
-#include "iree/compiler/Dialect/Stream/Transforms/PassDetail.h"
-#include "iree/compiler/Dialect/Stream/Transforms/Passes.h"
 #include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
 #include "iree/compiler/Dialect/Util/IR/UtilOps.h"
+#include "iree/compiler/Dialect/Util/Transforms/PassDetail.h"
+#include "iree/compiler/Dialect/Util/Transforms/Passes.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/Attributes.h"
@@ -23,7 +21,7 @@
 namespace mlir {
 namespace iree_compiler {
 namespace IREE {
-namespace Stream {
+namespace Util {
 
 // Returns true if |value| is worth outlining (large, etc).
 static bool isOutlinableValue(Attribute value) {
@@ -129,7 +127,7 @@ std::unique_ptr<OperationPass<mlir::ModuleOp>> createOutlineConstantsPass() {
   return std::make_unique<OutlineConstantsPass>();
 }
 
-} // namespace Stream
+} // namespace Util
 } // namespace IREE
 } // namespace iree_compiler
 } // namespace mlir
diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/Passes.h b/compiler/src/iree/compiler/Dialect/Util/Transforms/Passes.h
@@ -25,8 +25,11 @@ std::unique_ptr<OperationPass<mlir::ModuleOp>> createFoldGlobalsPass();
 std::unique_ptr<OperationPass<mlir::ModuleOp>> createFuseGlobalsPass();
 std::unique_ptr<OperationPass<mlir::ModuleOp>> createHoistIntoGlobalsPass();
 std::unique_ptr<OperationPass<mlir::ModuleOp>> createIPOPass();
+std::unique_ptr<OperationPass<mlir::ModuleOp>> createOutlineConstantsPass();
 std::unique_ptr<OperationPass<mlir::ModuleOp>> createPropagateSubrangesPass();
 std::unique_ptr<OperationPass<void>> createSimplifyGlobalAccessesPass();
+std::unique_ptr<OperationPass<mlir::ModuleOp>>
+createStripAndSplatConstantsPass();
 std::unique_ptr<OperationPass<void>> createStripDebugOpsPass();
 
 // Resource Management.

diff --git a/compiler/src/iree/compiler/Dialect/Util/Transforms/Passes.td b/compiler/src/iree/compiler/Dialect/Util/Transforms/Passes.td
@@ -54,13 +54,27 @@ def IPO : Pass<"iree-util-ipo", "mlir::ModuleOp"> {
   }];
 }
 
+def OutlineConstants :
+    Pass<"iree-util-outline-constants", "mlir::ModuleOp"> {
+  let summary = "Outlines tensor constants into util.globals at the module level.";
+  let constructor = [{
+    mlir::iree_compiler::IREE::Util::createOutlineConstantsPass()
+  }];
+}
+
 def PropagateSubranges : Pass<"iree-util-propagate-subranges", "mlir::ModuleOp"> {
   let summary = "Propagates resource subranges across the program.";
   let constructor = [{
     mlir::iree_compiler::IREE::Util::createPropagateSubrangesPass()
   }];
 }
 
+def StripAndSplatConstants :
+    Pass<"iree-util-strip-and-splat-constants", "mlir::ModuleOp"> {
+  let summary = "Strips constant util.global ops and replaces them with splats.";
+  let constructor = "mlir::iree_compiler::IREE::Util::createStripAndSplatConstantsPass()";
+}
+
 def StripDebugOps : Pass<"iree-util-strip-debug-ops", ""> {
   let summary = "Strips debug ops, like assertions.";
   let constructor = [{
@@ -112,7 +126,7 @@ def ImportResources : Pass<"iree-util-import-resources", ""> {
     derive from *ElementsAttr. Given the uniquing/inline behavior, this exacts
     very large runtime and memory overhead costs.
 
-    This is a temporary pass to convert a majority of the legacy 
+    This is a temporary pass to convert a majority of the legacy
     DenseElementsAttr attributes to DenseResourceElementsAttr. Ideally this
     is done at the source (frontend), but this pass is provided to aid
     transition and testing by doing a manual conversion with iree-opt.