llvm-project/mlir/lib/Dialect/Transform/Transforms/TransformInterpreterUtils.cpp

//===- TransformInterpreterUtils.cpp --------------------------------------===//
//
// Part of the LLVM Project, 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
//
//===----------------------------------------------------------------------===//
//
// Lightweight transform dialect interpreter utilities.
//
//===----------------------------------------------------------------------===//

#include "mlir/Dialect/Transform/Transforms/TransformInterpreterUtils.h"
#include "mlir/Dialect/Transform/IR/TransformDialect.h"
#include "mlir/Dialect/Transform/IR/TransformOps.h"
#include "mlir/Dialect/Transform/IR/Utils.h"
#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Verifier.h"
#include "mlir/IR/Visitors.h"
#include "mlir/Parser/Parser.h"
#include "mlir/Support/FileUtilities.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/raw_ostream.h"

using namespace mlir;

#define DEBUG_TYPE "transform-dialect-interpreter-utils"
#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")

/// Expands the given list of `paths` to a list of `.mlir` files.
///
/// Each entry in `paths` may either be a regular file, in which case it ends up
/// in the result list, or a directory, in which case all (regular) `.mlir`
/// files in that directory are added. Any other file types lead to a failure.
LogicalResult transform::detail::expandPathsToMLIRFiles(
    ArrayRef<std::string> paths, MLIRContext *context,
    SmallVectorImpl<std::string> &fileNames) {
  for (const std::string &path : paths) {
    auto loc = FileLineColLoc::get(context, path, 0, 0);

    if (llvm::sys::fs::is_regular_file(path)) {
      LLVM_DEBUG(DBGS() << "Adding '" << path << "' to list of files\n");
      fileNames.push_back(path);
      continue;
    }

    if (!llvm::sys::fs::is_directory(path)) {
      return emitError(loc)
             << "'" << path << "' is neither a file nor a directory";
    }

    LLVM_DEBUG(DBGS() << "Looking for files in '" << path << "':\n");

    std::error_code ec;
    for (llvm::sys::fs::directory_iterator it(path, ec), itEnd;
         it != itEnd && !ec; it.increment(ec)) {
      const std::string &fileName = it->path();

      if (it->type() != llvm::sys::fs::file_type::regular_file &&
          it->type() != llvm::sys::fs::file_type::symlink_file) {
        LLVM_DEBUG(DBGS() << "  Skipping non-regular file '" << fileName
                          << "'\n");
        continue;
      }

      if (!StringRef(fileName).ends_with(".mlir")) {
        LLVM_DEBUG(DBGS() << "  Skipping '" << fileName
                          << "' because it does not end with '.mlir'\n");
        continue;
      }

      LLVM_DEBUG(DBGS() << "  Adding '" << fileName << "' to list of files\n");
      fileNames.push_back(fileName);
    }

    if (ec)
      return emitError(loc) << "error while opening files in '" << path
                            << "': " << ec.message();
  }

  return success();
}

LogicalResult transform::detail::parseTransformModuleFromFile(
    MLIRContext *context, llvm::StringRef transformFileName,
    OwningOpRef<ModuleOp> &transformModule) {
  if (transformFileName.empty()) {
    LLVM_DEBUG(
        DBGS() << "no transform file name specified, assuming the transform "
                  "module is embedded in the IR next to the top-level\n");
    return success();
  }
  // Parse transformFileName content into a ModuleOp.
  std::string errorMessage;
  auto memoryBuffer = mlir::openInputFile(transformFileName, &errorMessage);
  if (!memoryBuffer) {
    return emitError(FileLineColLoc::get(
               StringAttr::get(context, transformFileName), 0, 0))
           << "failed to open transform file: " << errorMessage;
  }
  // Tell sourceMgr about this buffer, the parser will pick it up.
  llvm::SourceMgr sourceMgr;
  sourceMgr.AddNewSourceBuffer(std::move(memoryBuffer), llvm::SMLoc());
  transformModule =
      OwningOpRef<ModuleOp>(parseSourceFile<ModuleOp>(sourceMgr, context));
  if (!transformModule) {
    // Failed to parse the transform module.
    // Don't need to emit an error here as the parsing should have already done
    // that.
    return failure();
  }
  return mlir::verify(*transformModule);
}

ModuleOp transform::detail::getPreloadedTransformModule(MLIRContext *context) {
  return context->getOrLoadDialect<transform::TransformDialect>()
      ->getLibraryModule();
}

static transform::TransformOpInterface
findTransformEntryPointNonRecursive(Operation *op, StringRef entryPoint) {
  for (Region &region : op->getRegions()) {
    for (Block &block : region.getBlocks()) {
      for (auto namedSequenceOp : block.getOps<transform::NamedSequenceOp>()) {
        if (namedSequenceOp.getSymName() == entryPoint) {
          return cast<transform::TransformOpInterface>(
              namedSequenceOp.getOperation());
        }
      }
    }
  }
  return nullptr;
}

static transform::TransformOpInterface
findTransformEntryPointRecursive(Operation *op, StringRef entryPoint) {
  transform::TransformOpInterface transform = nullptr;
  op->walk<WalkOrder::PreOrder>(
      [&](transform::NamedSequenceOp namedSequenceOp) {
        if (namedSequenceOp.getSymName() == entryPoint) {
          transform = cast<transform::TransformOpInterface>(
              namedSequenceOp.getOperation());
          return WalkResult::interrupt();
        }
        return WalkResult::advance();
      });
  return transform;
}

// Will look for the transform's entry point favouring NamedSequenceOps
// ops that exist within the operation without the need for nesting.
// If no operation exists in the blocks owned by op, then it will recursively
// walk the op in preorder and find the first NamedSequenceOp that matches
// the entry point's name.
//
// This allows for the following two use cases:
// 1. op is a module annotated with the transform.with_named_sequence attribute
//    that has an entry point in its block. E.g.,
//
//    ```mlir
//    module {transform.with_named_sequence} {
//      transform.named_sequence @__transform_main(%arg0 : !transform.any_op) ->
//      () {
//        transform.yield
//      }
//    }
//    ```
//
// 2. op is a program which contains a nested module annotated with the
//    transform.with_named_sequence attribute. E.g.,
//
//    ```mlir
//    module {
//      func.func @foo () {
//      }
//
//      module {transform.with_named_sequence} {
//        transform.named_sequence @__transform_main(%arg0 : !transform.any_op)
//        -> () {
//          transform.yield
//        }
//      }
//    }
//    ```
static transform::TransformOpInterface
findTransformEntryPointInOp(Operation *op, StringRef entryPoint) {
  transform::TransformOpInterface transform =
      findTransformEntryPointNonRecursive(op, entryPoint);
  if (!transform)
    transform = findTransformEntryPointRecursive(op, entryPoint);
  return transform;
}

transform::TransformOpInterface
transform::detail::findTransformEntryPoint(Operation *root, ModuleOp module,
                                           StringRef entryPoint) {
  SmallVector<Operation *, 2> l{root};
  if (module)
    l.push_back(module);
  for (Operation *op : l) {
    TransformOpInterface transform =
        findTransformEntryPointInOp(op, entryPoint);
    if (transform)
      return transform;
  }
  auto diag = root->emitError()
              << "could not find a nested named sequence with name: "
              << entryPoint;
  return nullptr;
}

LogicalResult transform::detail::assembleTransformLibraryFromPaths(
    MLIRContext *context, ArrayRef<std::string> transformLibraryPaths,
    OwningOpRef<ModuleOp> &transformModule) {
  // Assemble list of library files.
  SmallVector<std::string> libraryFileNames;
  if (failed(detail::expandPathsToMLIRFiles(transformLibraryPaths, context,
                                            libraryFileNames)))
    return failure();

  // Parse modules from library files.
  SmallVector<OwningOpRef<ModuleOp>> parsedLibraries;
  for (const std::string &libraryFileName : libraryFileNames) {
    OwningOpRef<ModuleOp> parsedLibrary;
    if (failed(transform::detail::parseTransformModuleFromFile(
            context, libraryFileName, parsedLibrary)))
      return failure();
    parsedLibraries.push_back(std::move(parsedLibrary));
  }

  // Merge parsed libraries into one module.
  auto loc = FileLineColLoc::get(context, "<shared-library-module>", 0, 0);
  OwningOpRef<ModuleOp> mergedParsedLibraries =
      ModuleOp::create(loc, "__transform");
  {
    mergedParsedLibraries.get()->setAttr("transform.with_named_sequence",
                                         UnitAttr::get(context));
    // TODO: extend `mergeSymbolsInto` to support multiple `other` modules.
    for (OwningOpRef<ModuleOp> &parsedLibrary : parsedLibraries) {
      if (failed(transform::detail::mergeSymbolsInto(
              mergedParsedLibraries.get(), std::move(parsedLibrary))))
        return parsedLibrary->emitError()
               << "failed to merge symbols into shared library module";
    }
  }

  transformModule = std::move(mergedParsedLibraries);
  return success();
}

LogicalResult transform::applyTransformNamedSequence(
    Operation *payload, Operation *transformRoot, ModuleOp transformModule,
    const TransformOptions &options) {
  RaggedArray<MappedValue> bindings;
  bindings.push_back(ArrayRef<Operation *>{payload});
  return applyTransformNamedSequence(bindings,
                                     cast<TransformOpInterface>(transformRoot),
                                     transformModule, options);
}

LogicalResult transform::applyTransformNamedSequence(
    RaggedArray<MappedValue> bindings, TransformOpInterface transformRoot,
    ModuleOp transformModule, const TransformOptions &options) {
  if (bindings.empty()) {
    return transformRoot.emitError()
           << "expected at least one binding for the root";
  }
  if (bindings.at(0).size() != 1) {
    return transformRoot.emitError()
           << "expected one payload to be bound to the first argument, got "
           << bindings.at(0).size();
  }
  auto *payloadRoot = dyn_cast<Operation *>(bindings.at(0).front());
  if (!payloadRoot) {
    return transformRoot->emitError() << "expected the object bound to the "
                                         "first argument to be an operation";
  }

  bindings.removeFront();

  // `transformModule` may not be modified.
  if (transformModule && !transformModule->isAncestor(transformRoot)) {
    OwningOpRef<Operation *> clonedTransformModule(transformModule->clone());
    if (failed(detail::mergeSymbolsInto(
            SymbolTable::getNearestSymbolTable(transformRoot),
            std::move(clonedTransformModule)))) {
      return payloadRoot->emitError() << "failed to merge symbols";
    }
  }

  LLVM_DEBUG(DBGS() << "Apply\n" << *transformRoot << "\n");
  LLVM_DEBUG(DBGS() << "To\n" << *payloadRoot << "\n");

  return applyTransforms(payloadRoot, transformRoot, bindings, options,
                         /*enforceToplevelTransformOp=*/false);
}