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

// Copyright 2019 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/Dialect/VM/Transforms/Passes.h"

#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/Transforms/Passes.h"
#include "iree/compiler/Utils/PassUtils.h"
#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
#include "mlir/Dialect/Affine/Transforms/Passes.h"
#include "mlir/Dialect/Arith/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/Transforms/Passes.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Pass/PassRegistry.h"
#include "mlir/Transforms/Passes.h"

namespace mlir::iree_compiler::IREE::VM {

using FunctionLikeNest =
    MultiOpNest<func::FuncOp, IREE::Util::InitializerOp, IREE::Util::FuncOp>;

//===----------------------------------------------------------------------===//
// Utilities
//===----------------------------------------------------------------------===//

static void addCleanupPatterns(OpPassManager &passManager) {
  // TODO(benvanik): run in a fixed-point iteration pipeline.

  // Standard MLIR cleanup.
  FunctionLikeNest(passManager)
      .addPass(mlir::createCanonicalizerPass)
      .addPass(mlir::createCSEPass);

  // Aggressive MLIR cleanup.
  passManager.addNestedPass<IREE::VM::ModuleOp>(
      IREE::VM::createDropUnusedCallsPass());
  passManager.addPass(mlir::createSymbolDCEPass());

  FunctionLikeNest(passManager)
      // Simplify util.global accesses; this can help with data flow tracking as
      // redundant store-loads are removed.
      .addPass(IREE::Util::createSimplifyGlobalAccessesPass)

      // Aggressive cleanup.
      .addPass(IREE::Util::createApplyPatternsPass);

  // Cleanup and canonicalization of util.global (and other util ops).
  passManager.addPass(IREE::Util::createFoldGlobalsPass());
  passManager.addPass(IREE::Util::createFuseGlobalsPass());
}

static ConversionPassOptions
conversionPassOptionsFromTarget(IREE::VM::TargetOptions targetOptions) {
  ConversionPassOptions passOptions;
  passOptions.indexBits = targetOptions.indexBits;
  passOptions.f32Extension = targetOptions.f32Extension;
  passOptions.f64Extension = targetOptions.f64Extension;
  passOptions.truncateUnsupportedFloats =
      targetOptions.truncateUnsupportedFloats;
  passOptions.optimizeForStackSize = targetOptions.optimizeForStackSize;
  return passOptions;
}

//===----------------------------------------------------------------------===//
// -iree-vm-transformation-pipeline
//===----------------------------------------------------------------------===//

void buildVMTransformPassPipeline(OpPassManager &passManager,
                                  TargetOptions targetOptions) {
  // HACK: we'd really prefer not to inline here but most of the subsequent
  // passes are local only. We need to inline before we convert to the vm
  // dialect as they all operate on our various input dialects. If we could
  // control the inliner a bit more and not have it inline *everything* we could
  // do a more conservative inlining here and then the final one after lowering
  // to vm below.
  passManager.addPass(mlir::createInlinerPass());
  passManager.addPass(mlir::createSymbolDCEPass());

  // Verify module initialization order - subsequent passes rely on it being
  // correct (and we maintain it as correct from this point on, so this is our
  // gate).
  passManager.addPass(IREE::Util::createVerifyInitializationOrderPass());

  // Combine the initializers for all globals to allow us to optimize them
  // together.
  passManager.addPass(IREE::Util::createCombineInitializersPass());

  {
    MultiPipelineNest nest(passManager);
    nest.nest<func::FuncOp, IREE::Util::InitializerOp, IREE::Util::FuncOp>();

    nest.addPass(mlir::createSCFForLoopCanonicalizationPass);

    // This pass is sketchy as it can pessimize tight loops due to affine
    // treating all indices as signed and the unsigned conversion pass not being
    // able to handle that. The scf.for canonicalization does a decent job of
    // removing trivial loops above and this catches the rest. It inserts nasty
    // rem/div ops that we can never safely remove inside of the hot inner loop
    // and that sucks. We still have this here for now as the cost of the
    // rem/div are less than the cost of an additional loop that this could
    // remove.
    nest.addPassFor<func::FuncOp>(affine::createLoopCoalescingPass);

    nest.addPass(mlir::createLoopInvariantCodeMotionPass)
        .addPass(mlir::createSCFToControlFlowPass)
        // TODO: Maybe this should be a part of Affine lowering pass.
        // Remove if it is added there.
        // https://github.com/llvm/llvm-project/issues/78458
        .addPass(createIREEAffineExpandIndexOpsPass)
        .addPass(createIREELowerAffinePass);
  }

  passManager.addPass(mlir::arith::createArithUnsignedWhenEquivalentPass());

  // 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(conversionPassOptionsFromTarget(targetOptions)));

  // Hoist globals and get the final set that need to be initialized.
  passManager.addNestedPass<IREE::VM::ModuleOp>(createReifyRodataTablesPass());
  passManager.addNestedPass<IREE::VM::ModuleOp>(createHoistInlinedRodataPass());
  passManager.addNestedPass<IREE::VM::ModuleOp>(createDeduplicateRodataPass());
  addCleanupPatterns(passManager);
  passManager.addNestedPass<IREE::VM::ModuleOp>(createResolveRodataLoadsPass());

  // Catch any inlining opportunities we created during lowering.
  passManager.addPass(mlir::createInlinerPass());
  passManager.addPass(mlir::createSymbolDCEPass());

  // Create global initialization functions (__init/__deinit).
  addCleanupPatterns(passManager);
  passManager.addNestedPass<IREE::VM::ModuleOp>(
      createGlobalInitializationPass());

  // Ideally we'd run this as part of a fixed-point iteration: CSE may need the
  // canonicalizer to remove ops in order to get the equivalences it needs to
  // work while the canonicalizer may require CSE for patterns to kick in (like
  // "fold cmp if lhs=%0 && rhs=%0").
  passManager.addPass(mlir::createCanonicalizerPass());
  passManager.addPass(mlir::createCSEPass());
  passManager.addPass(mlir::createCanonicalizerPass());

  // Now that we've inlined/canonicalized/etc the initializers we can remove
  // them if they are empty to save a few bytes in the binary and avoid a
  // runtime initialization call.
  passManager.addNestedPass<IREE::VM::ModuleOp>(
      createDropEmptyModuleInitializersPass());

  // Annotate functions with yield/unwind requirements by analyzing the call
  // graph. This must happen after all inlining is complete.
  passManager.addNestedPass<IREE::VM::ModuleOp>(createAnnotateFunctionsPass());

  // Convert calls to yieldable functions to vm.call.yieldable. This must run
  // after AnnotateFunctionsPass which marks functions with vm.yield.
  passManager.addNestedPass<IREE::VM::ModuleOp>(
      createConvertToYieldableCallsPass());

  if (targetOptions.optimizeForStackSize) {
    passManager.addNestedPass<IREE::VM::ModuleOp>(createSinkDefiningOpsPass());
  }

  // Drop vm.optimization_barrier ops now that optimization is complete.
  passManager.addNestedPass<IREE::VM::ModuleOp>(
      createDropOptimizationBarriersPass());

  // Insert explicit discard ops for ref values at their last use points.
  // Uses edge-based placement: refs dying on control flow edges get discards
  // inserted on those edges, refs dying mid-block get discards after last use.
  passManager.addNestedPass<IREE::VM::ModuleOp>(
      createMaterializeRefDiscardsPass());
}

//===----------------------------------------------------------------------===//
// Registration
//===----------------------------------------------------------------------===//

namespace {
#define GEN_PASS_REGISTRATION
#include "iree/compiler/Dialect/VM/Transforms/Passes.h.inc" // IWYU pragma: export
} // namespace

void registerVMPasses() {
  // Register command line flags.
  IREE::VM::TargetOptions::FromFlags::get();

  // Generated.
  registerPasses();

  // Pipelines.
  PassPipelineRegistration<> transformPassPipeline(
      "iree-vm-transformation-pipeline",
      "Runs the full IREE VM dialect transformation pipeline",
      [](OpPassManager &passManager) {
        buildVMTransformPassPipeline(passManager,
                                     IREE::VM::TargetOptions::FromFlags::get());
      });
}

} // namespace mlir::iree_compiler::IREE::VM