[mlir][SMT] add export smtlib (llvm#131492)

This PR adds the `ExportSMTLIB` translation/egress pass for `SMT`
dialect.

Author: makslevental

mlir/include/mlir/InitAllTranslations.h

@@ -22,6 +22,10 @@ void registerToCppTranslation();
 void registerToLLVMIRTranslation();
 void registerToSPIRVTranslation();
 
+namespace smt {
+void registerExportSMTLIBTranslation();
+}
+
 // This function should be called before creating any MLIRContext if one
 // expects all the possible translations to be made available to the context
 // automatically.
@@ -32,6 +36,7 @@ inline void registerAllTranslations() {
     registerToCppTranslation();
     registerToLLVMIRTranslation();
     registerToSPIRVTranslation();
+    smt::registerExportSMTLIBTranslation();
     return true;
   }();
   (void)initOnce;

mlir/include/mlir/Target/SMTLIB/ExportSMTLIB.h

@@ -0,0 +1,43 @@
+//===- ExportSMTLIB.h - SMT-LIB Exporter ------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines the interface to the SMT-LIB emitter.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TARGET_EXPORTSMTLIB_H
+#define MLIR_TARGET_EXPORTSMTLIB_H
+
+#include "mlir/Support/LLVM.h"
+
+namespace mlir {
+class Operation;
+namespace smt {
+
+/// Emission options for the ExportSMTLIB pass. Allows controlling the emitted
+/// format and overall behavior.
+struct SMTEmissionOptions {
+  // Don't produce 'let' expressions to bind expressions that are only used
+  // once, but inline them directly at the use-site.
+  bool inlineSingleUseValues = false;
+  // Increase indentation for each 'let' expression body.
+  bool indentLetBody = false;
+};
+
+/// Run the ExportSMTLIB pass.
+LogicalResult
+exportSMTLIB(Operation *module, llvm::raw_ostream &os,
+             const SMTEmissionOptions &options = SMTEmissionOptions());
+
+/// Register the ExportSMTLIB pass.
+void registerExportSMTLIBTranslation();
+
+} // namespace smt
+} // namespace mlir
+
+#endif // MLIR_TARGET_EXPORTSMTLIB_H

mlir/include/mlir/Target/SMTLIB/Namespace.h

@@ -0,0 +1,170 @@
+//===- Namespace.h - Utilities for generating names -------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file provides utilities for generating new names that do not conflict
+// with existing names.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_SUPPORT_NAMESPACE_H
+#define MLIR_SUPPORT_NAMESPACE_H
+
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Target/SMTLIB/SymCache.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringSet.h"
+#include "llvm/ADT/Twine.h"
+
+namespace mlir {
+
+/// A namespace that is used to store existing names and generate new names in
+/// some scope within the IR. This exists to work around limitations of
+/// SymbolTables. This acts as a base class providing facilities common to all
+/// namespaces implementations.
+class Namespace {
+public:
+  Namespace() {
+    // This fills an entry for an empty string beforehand so that `newName`
+    // doesn't return an empty string.
+    nextIndex.insert({"", 0});
+  }
+  Namespace(const Namespace &other) = default;
+  Namespace(Namespace &&other)
+      : nextIndex(std::move(other.nextIndex)), locked(other.locked) {}
+
+  Namespace &operator=(const Namespace &other) = default;
+  Namespace &operator=(Namespace &&other) {
+    nextIndex = std::move(other.nextIndex);
+    locked = other.locked;
+    return *this;
+  }
+
+  void add(mlir::ModuleOp module) {
+    assert(module->getNumRegions() == 1);
+    for (auto &op : module.getBody(0)->getOperations())
+      if (auto symbol = op.getAttrOfType<mlir::StringAttr>(
+              mlir::SymbolTable::getSymbolAttrName()))
+        nextIndex.insert({symbol.getValue(), 0});
+  }
+
+  /// SymbolCache initializer; initialize from every key that is convertible to
+  /// a StringAttr in the SymbolCache.
+  void add(SymbolCache &symCache) {
+    for (auto &&[attr, _] : symCache)
+      if (auto strAttr = dyn_cast<mlir::StringAttr>(attr))
+        nextIndex.insert({strAttr.getValue(), 0});
+  }
+
+  void add(llvm::StringRef name) { nextIndex.insert({name, 0}); }
+
+  /// Removes a symbol from the namespace. Returns true if the symbol was
+  /// removed, false if the symbol was not found.
+  /// This is only allowed to be called _before_ any call to newName.
+  bool erase(llvm::StringRef symbol) {
+    assert(!locked && "Cannot erase names from a locked namespace");
+    return nextIndex.erase(symbol);
+  }
+
+  /// Empty the namespace.
+  void clear() {
+    nextIndex.clear();
+    locked = false;
+  }
+
+  /// Return a unique name, derived from the input `name`, and add the new name
+  /// to the internal namespace.  There are two possible outcomes for the
+  /// returned name:
+  ///
+  /// 1. The original name is returned.
+  /// 2. The name is given a `_<n>` suffix where `<n>` is a number starting from
+  ///    `0` and incrementing by one each time (`_0`, ...).
+  llvm::StringRef newName(const llvm::Twine &name) {
+    locked = true;
+    // Special case the situation where there is no name collision to avoid
+    // messing with the SmallString allocation below.
+    llvm::SmallString<64> tryName;
+    auto inserted = nextIndex.insert({name.toStringRef(tryName), 0});
+    if (inserted.second)
+      return inserted.first->getKey();
+
+    // Try different suffixes until we get a collision-free one.
+    if (tryName.empty())
+      name.toVector(tryName); // toStringRef may leave tryName unfilled
+
+    // Indexes less than nextIndex[tryName] are lready used, so skip them.
+    // Indexes larger than nextIndex[tryName] may be used in another name.
+    size_t &i = nextIndex[tryName];
+    tryName.push_back('_');
+    size_t baseLength = tryName.size();
+    do {
+      tryName.resize(baseLength);
+      llvm::Twine(i++).toVector(tryName); // append integer to tryName
+      inserted = nextIndex.insert({tryName, 0});
+    } while (!inserted.second);
+
+    return inserted.first->getKey();
+  }
+
+  /// Return a unique name, derived from the input `name` and ensure the
+  /// returned name has the input `suffix`. Also add the new name to the
+  /// internal namespace.
+  /// There are two possible outcomes for the returned name:
+  /// 1. The original name + `_<suffix>` is returned.
+  /// 2. The name is given a suffix `_<n>_<suffix>` where `<n>` is a number
+  ///    starting from `0` and incrementing by one each time.
+  llvm::StringRef newName(const llvm::Twine &name, const llvm::Twine &suffix) {
+    locked = true;
+    // Special case the situation where there is no name collision to avoid
+    // messing with the SmallString allocation below.
+    llvm::SmallString<64> tryName;
+    auto inserted = nextIndex.insert(
+        {name.concat("_").concat(suffix).toStringRef(tryName), 0});
+    if (inserted.second)
+      return inserted.first->getKey();
+
+    // Try different suffixes until we get a collision-free one.
+    tryName.clear();
+    name.toVector(tryName); // toStringRef may leave tryName unfilled
+    tryName.push_back('_');
+    size_t baseLength = tryName.size();
+
+    // Get the initial number to start from.  Since `:` is not a valid character
+    // in a verilog identifier, we use it separate the name and suffix.
+    // Next number for name+suffix is stored with key `name_:suffix`.
+    tryName.push_back(':');
+    suffix.toVector(tryName);
+
+    // Indexes less than nextIndex[tryName] are already used, so skip them.
+    // Indexes larger than nextIndex[tryName] may be used in another name.
+    size_t &i = nextIndex[tryName];
+    do {
+      tryName.resize(baseLength);
+      llvm::Twine(i++).toVector(tryName); // append integer to tryName
+      tryName.push_back('_');
+      suffix.toVector(tryName);
+      inserted = nextIndex.insert({tryName, 0});
+    } while (!inserted.second);
+
+    return inserted.first->getKey();
+  }
+
+protected:
+  // The "next index" that will be tried when trying to unique a string within a
+  // namespace.  It follows that all values less than the "next index" value are
+  // already used.
+  llvm::StringMap<size_t> nextIndex;
+
+  // When true, no names can be erased from the namespace. This is to prevent
+  // erasing names after they have been used, thus leaving users of the
+  // namespace in an inconsistent state.
+  bool locked = false;
+};
+
+} // namespace mlir
+
+#endif // MLIR_SUPPORT_NAMESPACE_H

mlir/include/mlir/Target/SMTLIB/SymCache.h

@@ -0,0 +1,133 @@
+//===- SymCache.h - Declare Symbol Cache ------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares a Symbol Cache.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_SUPPORT_SYMCACHE_H
+#define MLIR_SUPPORT_SYMCACHE_H
+
+#include "mlir/IR/SymbolTable.h"
+#include "llvm/ADT/iterator.h"
+#include "llvm/Support/Casting.h"
+
+namespace mlir {
+
+/// Base symbol cache class to allow for cache lookup through a pointer to some
+/// abstract cache. A symbol cache stores lookup tables to make manipulating and
+/// working with the IR more efficient.
+class SymbolCacheBase {
+public:
+  virtual ~SymbolCacheBase();
+
+  /// Defines 'op' as associated with the 'symbol' in the cache.
+  virtual void addDefinition(mlir::Attribute symbol, mlir::Operation *op) = 0;
+
+  /// Adds the symbol-defining 'op' to the cache.
+  void addSymbol(mlir::SymbolOpInterface op) {
+    addDefinition(op.getNameAttr(), op);
+  }
+
+  /// Populate the symbol cache with all symbol-defining operations within the
+  /// 'top' operation.
+  void addDefinitions(mlir::Operation *top);
+
+  /// Lookup a definition for 'symbol' in the cache.
+  virtual mlir::Operation *getDefinition(mlir::Attribute symbol) const = 0;
+
+  /// Lookup a definition for 'symbol' in the cache.
+  mlir::Operation *getDefinition(mlir::FlatSymbolRefAttr symbol) const {
+    return getDefinition(symbol.getAttr());
+  }
+
+  /// Iterator support through a pointer to some abstract cache.
+  /// The implementing cache must provide an iterator that carries values on the
+  /// form of <mlir::Attribute, mlir::Operation*>.
+  using CacheItem = std::pair<mlir::Attribute, mlir::Operation *>;
+  struct CacheIteratorImpl {
+    virtual ~CacheIteratorImpl() {}
+    virtual void operator++() = 0;
+    virtual CacheItem operator*() = 0;
+    virtual bool operator==(CacheIteratorImpl *other) = 0;
+  };
+
+  struct Iterator
+      : public llvm::iterator_facade_base<Iterator, std::forward_iterator_tag,
+                                          CacheItem> {
+    Iterator(std::unique_ptr<CacheIteratorImpl> &&impl)
+        : impl(std::move(impl)) {}
+    CacheItem operator*() const { return **impl; }
+    using llvm::iterator_facade_base<Iterator, std::forward_iterator_tag,
+                                     CacheItem>::operator++;
+    bool operator==(const Iterator &other) const {
+      return *impl == other.impl.get();
+    }
+    void operator++() { impl->operator++(); }
+
+  private:
+    std::unique_ptr<CacheIteratorImpl> impl;
+  };
+  virtual Iterator begin() = 0;
+  virtual Iterator end() = 0;
+};
+
+/// Default symbol cache implementation; stores associations between names
+/// (StringAttr's) to mlir::Operation's.
+/// Adding/getting definitions from the symbol cache is not
+/// thread safe. If this is required, synchronizing cache acccess should be
+/// ensured by the caller.
+class SymbolCache : public SymbolCacheBase {
+public:
+  /// In the building phase, add symbols.
+  void addDefinition(mlir::Attribute key, mlir::Operation *op) override {
+    symbolCache.try_emplace(key, op);
+  }
+
+  // Pull in getDefinition(mlir::FlatSymbolRefAttr symbol)
+  using SymbolCacheBase::getDefinition;
+  mlir::Operation *getDefinition(mlir::Attribute attr) const override {
+    auto it = symbolCache.find(attr);
+    if (it == symbolCache.end())
+      return nullptr;
+    return it->second;
+  }
+
+protected:
+  /// This stores a lookup table from symbol attribute to the operation
+  /// that defines it.
+  llvm::DenseMap<mlir::Attribute, mlir::Operation *> symbolCache;
+
+private:
+  /// Iterator support: A simple mapping between decltype(symbolCache)::iterator
+  /// to SymbolCacheBase::Iterator.
+  using Iterator = decltype(symbolCache)::iterator;
+  struct SymbolCacheIteratorImpl : public CacheIteratorImpl {
+    SymbolCacheIteratorImpl(Iterator it) : it(it) {}
+    CacheItem operator*() override { return {it->getFirst(), it->getSecond()}; }
+    void operator++() override { it++; }
+    bool operator==(CacheIteratorImpl *other) override {
+      return it == static_cast<SymbolCacheIteratorImpl *>(other)->it;
+    }
+    Iterator it;
+  };
+
+public:
+  SymbolCacheBase::Iterator begin() override {
+    return SymbolCacheBase::Iterator(
+        std::make_unique<SymbolCacheIteratorImpl>(symbolCache.begin()));
+  }
+  SymbolCacheBase::Iterator end() override {
+    return SymbolCacheBase::Iterator(
+        std::make_unique<SymbolCacheIteratorImpl>(symbolCache.end()));
+  }
+};
+
+} // namespace mlir
+
+#endif // MLIR_SUPPORT_SYMCACHE_H

mlir/lib/Target/CMakeLists.txt

@@ -2,3 +2,4 @@ add_subdirectory(Cpp)
 add_subdirectory(SPIRV)
 add_subdirectory(LLVMIR)
 add_subdirectory(LLVM)
+add_subdirectory(SMTLIB)

mlir/lib/Target/SMTLIB/CMakeLists.txt

@@ -0,0 +1,13 @@
+add_mlir_translation_library(MLIRExportSMTLIB
+  ExportSMTLIB.cpp
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRSMT
+  MLIRSupport
+  MLIRFuncDialect
+  MLIRIR
+  MLIRTranslateLib
+)