Running buffer subspan propagation during vm transformation.

compiler/src/iree/compiler/Dialect/VM/Transforms/BUILD

diff --git a/compiler/src/iree/compiler/Dialect/VM/Transforms/BUILD b/compiler/src/iree/compiler/Dialect/VM/Transforms/BUILD
index d3eb105..ba3cced 100644
--- a/compiler/src/iree/compiler/Dialect/VM/Transforms/BUILD
+++ b/compiler/src/iree/compiler/Dialect/VM/Transforms/BUILD
@@ -30,6 +30,7 @@
     deps = [
         "//compiler/src/iree/compiler/Dialect/Util/Conversion",
         "//compiler/src/iree/compiler/Dialect/Util/IR",
+        "//compiler/src/iree/compiler/Dialect/Util/Transforms",
         "//compiler/src/iree/compiler/Dialect/VM/Conversion",
         "//compiler/src/iree/compiler/Dialect/VM/Conversion/MathToVM",
         "//compiler/src/iree/compiler/Dialect/VM/Conversion/MemRefToVM",

compiler/src/iree/compiler/Dialect/VM/Transforms/CMakeLists.txt

diff --git a/compiler/src/iree/compiler/Dialect/VM/Transforms/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/VM/Transforms/CMakeLists.txt
index 96b1ad8..c51c863 100644
--- a/compiler/src/iree/compiler/Dialect/VM/Transforms/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Dialect/VM/Transforms/CMakeLists.txt
@@ -44,6 +44,7 @@
     MLIRTransforms
     iree::compiler::Dialect::Util::Conversion
     iree::compiler::Dialect::Util::IR
+    iree::compiler::Dialect::Util::Transforms
     iree::compiler::Dialect::VM::Conversion
     iree::compiler::Dialect::VM::Conversion::MathToVM
     iree::compiler::Dialect::VM::Conversion::MemRefToVM

compiler/src/iree/compiler/Dialect/VM/Transforms/Conversion.cpp

diff --git a/compiler/src/iree/compiler/Dialect/VM/Transforms/Conversion.cpp b/compiler/src/iree/compiler/Dialect/VM/Transforms/Conversion.cpp
index 7e2d4f8..4df294d 100644
--- a/compiler/src/iree/compiler/Dialect/VM/Transforms/Conversion.cpp
+++ b/compiler/src/iree/compiler/Dialect/VM/Transforms/Conversion.cpp
@@ -69,18 +69,18 @@
   explicit ConversionPass(TargetOptions targetOptions)
       : targetOptions_(targetOptions) {}
 
-  void getDependentDialects(DialectRegistry &registry) const override {
-    registry.insert<IREE::Util::UtilDialect, IREE::VM::VMDialect,
-                    func::FuncDialect, mlir::arith::ArithmeticDialect,
-                    math::MathDialect, AffineDialect, memref::MemRefDialect>();
-  }
-
   StringRef getArgument() const override { return "iree-vm-conversion"; }
 
   StringRef getDescription() const override {
     return "Converts from various dialects to the VM dialect";
   }
 
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<IREE::Util::UtilDialect, IREE::VM::VMDialect,
+                    func::FuncDialect, mlir::arith::ArithmeticDialect,
+                    math::MathDialect, AffineDialect, memref::MemRefDialect>();
+  }
+
   void runOnOperation() override {
     if (getOperation().getBody()->empty()) return;
 
@@ -116,17 +116,17 @@
       }
     }
 
-    RewritePatternSet conversionPatterns(&getContext());
+    RewritePatternSet patterns(&getContext());
     populateUtilConversionPatterns(context, conversionTarget, typeConverter,
-                                   conversionPatterns);
+                                   patterns);
     populateUtilToVMPatterns(context, conversionTarget, typeConverter,
-                             conversionPatterns);
-    arith::populateArithmeticExpandOpsPatterns(conversionPatterns);
-    populateStandardToVMPatterns(context, typeConverter, conversionPatterns);
-    populateMathToVMPatterns(context, typeConverter, conversionPatterns);
+                             patterns);
+    arith::populateArithmeticExpandOpsPatterns(patterns);
+    populateStandardToVMPatterns(context, typeConverter, patterns);
+    populateMathToVMPatterns(context, typeConverter, patterns);
     populateMemRefToVMPatterns(context, conversionTarget, typeConverter,
-                               conversionPatterns);
-    populateAffineToStdConversionPatterns(conversionPatterns);
+                               patterns);
+    populateAffineToStdConversionPatterns(patterns);
 
     conversionTarget
         .addIllegalDialect<func::FuncDialect, mlir::arith::ArithmeticDialect>();
@@ -138,11 +138,11 @@
     SymbolTable importSymbols(innerModuleOp);
     for (auto *dialectInterface : usedDialects) {
       dialectInterface->populateVMConversionPatterns(
-          importSymbols, conversionPatterns, conversionTarget, typeConverter);
+          importSymbols, patterns, conversionTarget, typeConverter);
     }
 
     if (failed(applyPartialConversion(outerModuleOp, conversionTarget,
-                                      std::move(conversionPatterns)))) {
+                                      std::move(patterns)))) {
       outerModuleOp.emitError() << "conversion to vm.module failed";
       return signalPassFailure();
     }

compiler/src/iree/compiler/Dialect/VM/Transforms/Passes.cpp

diff --git a/compiler/src/iree/compiler/Dialect/VM/Transforms/Passes.cpp b/compiler/src/iree/compiler/Dialect/VM/Transforms/Passes.cpp
index 6d8e8d7..db45bdf 100644
--- a/compiler/src/iree/compiler/Dialect/VM/Transforms/Passes.cpp
+++ b/compiler/src/iree/compiler/Dialect/VM/Transforms/Passes.cpp
@@ -9,6 +9,7 @@
 #include <memory>
 
 #include "iree/compiler/Dialect/Util/IR/UtilOps.h"
+#include "iree/compiler/Dialect/Util/Transforms/Passes.h"
 #include "iree/compiler/Dialect/VM/IR/VMOps.h"
 #include "iree/compiler/Utils/PassUtils.h"
 #include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
@@ -26,6 +27,32 @@
 
 using FunctionLikeNest = MultiOpNest<func::FuncOp, IREE::Util::InitializerOp>;
 
+//===----------------------------------------------------------------------===//
+// Utilities
+//===----------------------------------------------------------------------===//
+
+static void addCleanupPatterns(OpPassManager &passManager) {
+  // TODO(benvanik): run in a fixed-point iteration pipeline.
+
+  // Standard MLIR cleanup.
+  passManager.addPass(mlir::createCanonicalizerPass());
+  passManager.addPass(mlir::createCSEPass());
+
+  // Simplify util.global accesses; this can help with data flow tracking as
+  // redundant store-loads are removed.
+  FunctionLikeNest(passManager)
+      .addPass(IREE::Util::createSimplifyGlobalAccessesPass);
+
+  // Cleanup and canonicalization of util.global (and other util ops).
+  passManager.addPass(IREE::Util::createApplyPatternsPass());
+  passManager.addPass(IREE::Util::createFoldGlobalsPass());
+  passManager.addPass(IREE::Util::createFuseGlobalsPass());
+}
+
+//===----------------------------------------------------------------------===//
+// -iree-vm-transformation-pipeline
+//===----------------------------------------------------------------------===//
+
 void buildVMTransformPassPipeline(OpPassManager &passManager,
                                   TargetOptions targetOptions) {
   passManager.addNestedPass<mlir::func::FuncOp>(createLoopCoalescingPass());
@@ -35,9 +62,13 @@
       .addPass(createLoopInvariantCodeMotionPass)
       .addPass(createConvertSCFToCFPass);
 
-  passManager.addPass(createCanonicalizerPass());
-  passManager.addPass(createCSEPass());
+  // Propagate buffer subranges throughout the program - this should remove any
+  // remaining subspans and give us a smaller surface area during conversion.
+  passManager.addPass(IREE::Util::createPropagateSubrangesPass());
+  addCleanupPatterns(passManager);
 
+  // Convert std/util/etc -> VM, along with any other dialects implementing the
+  // VM conversion dialect interface.
   passManager.addPass(createConversionPass(targetOptions));
 
   passManager.addNestedPass<IREE::VM::ModuleOp>(createHoistInlinedRodataPass());