torch-mlir/python/npcomp/mlir_edsc.cpp

//===- mlir_edsc.cpp - MLIR EDSC bindings ---------------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//

#include <cstddef>
#include <unordered_map>

#include <pybind11/pybind11.h>
#include <pybind11/pytypes.h>
#include <pybind11/stl.h>

#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir-c/Core.h"
#include "mlir/Conversion/LoopToStandard/ConvertLoopToStandard.h"
#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
#include "mlir/Dialect/Affine/EDSC/Intrinsics.h"
#include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
#include "mlir/EDSC/Builders.h"
#include "mlir/EDSC/Intrinsics.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Function.h"
#include "mlir/IR/Module.h"
#include "mlir/IR/Types.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Target/LLVMIR.h"
#include "mlir/Transforms/Passes.h"

namespace py = pybind11;
using mlir::edsc::intrinsics::AffineIndexedValue;

using namespace mlir;
using namespace mlir::edsc;

namespace mlir {
namespace edsc {

struct CustomValueBuilder {
  CustomValueBuilder(StringRef name, ArrayRef<Value> operands,
                     ArrayRef<Type> resultTypes,
                     ArrayRef<NamedAttribute> attributes = {}) {
    OperationState state(ScopedContext::getLocation(), name);
    SmallVector<Value, 4> ops(operands.begin(), operands.end());
    state.addOperands(ops);
    state.addTypes(resultTypes);
    for (const auto &attr : attributes)
      state.addAttribute(attr.first, attr.second);
    Operation *op = ScopedContext::getBuilder().createOperation(state);
    if (op->getNumResults() != 1)
      llvm_unreachable("unsupported operation, use an OperationHandle instead");
    value = op->getResult(0);
  }
  operator Value() { return value; }
  Value value;
};

struct PythonAttribute;
struct PythonAttributedType;
struct PythonBindable;
struct PythonExpr;
struct PythonFunctionContext;
struct PythonStmt;
struct PythonBlock;
struct PythonAffineExpr;
struct PythonAffineMap;

struct PythonType {
  PythonType() : type{nullptr} {}
  PythonType(mlir_type_t t) : type{t} {}

  operator mlir_type_t() const { return type; }

  PythonAttributedType attachAttributeDict(
      const std::unordered_map<std::string, PythonAttribute> &attrs) const;

  std::string str() {
    mlir::Type f = mlir::Type::getFromOpaquePointer(type);
    std::string res;
    llvm::raw_string_ostream os(res);
    f.print(os);
    return res;
  }

  mlir_type_t type;
};

struct PythonValueHandle {
  PythonValueHandle(const PythonValueHandle &other) = default;
  PythonValueHandle(const mlir::Value &other) : value(other) {}
  PythonValueHandle(mlir::edsc::CustomValueBuilder vb) : value(vb) {}
  template <typename Op>
  PythonValueHandle(mlir::edsc::ValueBuilder<Op> vb) : value(vb) {}
  operator mlir::Value() const { return value; }
  operator mlir::Value &() { return value; }

  std::string str() const {
    return std::to_string(
        reinterpret_cast<intptr_t>(value.getAsOpaquePointer()));
  }

  PythonValueHandle call(const std::vector<PythonValueHandle> &args) {
    std::vector<Value> argValues;
    argValues.reserve(args.size());
    for (auto arg : args) argValues.push_back(arg.value);
    return mlir::edsc::ValueBuilder<CallIndirectOp>(value, argValues);
  }

  PythonType type() const {
    return PythonType(value.getType().getAsOpaquePointer());
  }

  mlir::Value value;
};

struct PythonFunction {
  PythonFunction() : function{nullptr} {}
  PythonFunction(mlir_func_t f) : function{f} {}
  PythonFunction(mlir::FuncOp f)
      : function(const_cast<void *>(f.getAsOpaquePointer())) {}
  operator mlir_func_t() { return function; }
  std::string str() {
    mlir::FuncOp f = mlir::FuncOp::getFromOpaquePointer(function);
    std::string res;
    llvm::raw_string_ostream os(res);
    f.print(os);
    return res;
  }

  // If the function does not yet have an entry block, i.e. if it is a function
  // declaration, add the entry block, transforming the declaration into a
  // definition.  Return true if the block was added, false otherwise.
  bool define() {
    auto f = mlir::FuncOp::getFromOpaquePointer(function);
    if (!f.getBlocks().empty()) return false;

    f.addEntryBlock();
    return true;
  }

  PythonValueHandle arg(unsigned index) {
    auto f = mlir::FuncOp::getFromOpaquePointer(function);
    assert(index < f.getNumArguments() && "argument index out of bounds");
    return PythonValueHandle(f.getArgument(index));
  }

  mlir_func_t function;
};

/// Trivial C++ wrappers make use of the EDSC C API.
struct PythonMLIRModule {
  PythonMLIRModule()
      : mlirContext(),
        module(mlir::ModuleOp::create(mlir::UnknownLoc::get(&mlirContext))),
        symbolTable(*module) {}

  PythonType makeMemRefType(PythonType elemType, std::vector<int64_t> sizes) {
    return ::makeMemRefType(mlir_context_t{&mlirContext}, elemType,
                            int64_list_t{sizes.data(), sizes.size()});
  }
  PythonType makeIndexType() {
    return ::makeIndexType(mlir_context_t{&mlirContext});
  }
  PythonType makeType(const std::string &type) {
    return ::mlirParseType(type.c_str(), mlir_context_t{&mlirContext}, nullptr);
  }

  // Declare a function with the given name, input types and their attributes,
  // output types, and function attributes, but do not define it.
  PythonFunction declareFunction(const std::string &name,
                                 const py::list &inputs,
                                 const std::vector<PythonType> &outputTypes,
                                 const py::kwargs &funcAttributes);

  // Declare a function with the given name, input types and their attributes,
  // output types, and function attributes.
  PythonFunction makeFunction(const std::string &name, const py::list &inputs,
                              const std::vector<PythonType> &outputTypes,
                              const py::kwargs &funcAttributes) {
    auto declaration =
        declareFunction(name, inputs, outputTypes, funcAttributes);
    declaration.define();
    return declaration;
  }

  // Create a custom op given its name and arguments.
  PythonExpr op(const std::string &name, PythonType type,
                const py::list &arguments, const py::list &successors,
                py::kwargs attributes);

  // Creates an integer attribute.
  PythonAttribute integerAttr(PythonType type, int64_t value);

  // Creates a boolean attribute.
  PythonAttribute boolAttr(bool value);

  // Creates a float attribute.
  PythonAttribute floatAttr(float value);

  // Creates a string atrribute.
  PythonAttribute stringAttr(const std::string &value);

  // Creates an Array attribute.
  PythonAttribute arrayAttr(const std::vector<PythonAttribute> &values);

  // Creates an AffineMap attribute.
  PythonAttribute affineMapAttr(PythonAffineMap value);

  // Creates a dictionary attribute.
  PythonAttribute dictionaryAttr(
      const std::unordered_map<std::string, PythonAttribute> &attrs);

  // Creates an affine constant expression.
  PythonAffineExpr affineConstantExpr(int64_t value);

  // Creates an affine symbol expression.
  PythonAffineExpr affineSymbolExpr(unsigned position);

  // Creates an affine dimension expression.
  PythonAffineExpr affineDimExpr(unsigned position);

  // Creates a single constant result affine map.
  PythonAffineMap affineConstantMap(int64_t value);

  // Creates an affine map.
  PythonAffineMap affineMap(unsigned dimCount, unsigned symbolCount,
                            const std::vector<PythonAffineExpr> &results);

  std::string getIR() {
    std::string res;
    llvm::raw_string_ostream os(res);
    module->print(os);
    return res;
  }

  PythonFunction getNamedFunction(const std::string &name) {
    return symbolTable.lookup<FuncOp>(name);
  }

  PythonFunctionContext makeNewFunctionContext(
      const std::string &name, const py::list &inputs,
      const std::vector<PythonType> &outputs, const py::kwargs &attributes);

  PythonFunctionContext makeFunctionContext(
      PythonFunction f);

 private:
  mlir::MLIRContext mlirContext;
  // One single module in a python-exposed MLIRContext for now.
  mlir::OwningModuleRef module;
  mlir::SymbolTable symbolTable;
};

struct PythonFunctionContext {
  PythonFunctionContext(PythonFunction f) : function(f) {}
  PythonFunctionContext(PythonMLIRModule &module, const std::string &name,
                        const py::list &inputs,
                        const std::vector<PythonType> &outputs,
                        const py::kwargs &attributes) {
    auto function = module.declareFunction(name, inputs, outputs, attributes);
    function.define();
  }

  PythonFunction enter() {
    assert(function.function && "function is not set up");
    auto mlirFunc = mlir::FuncOp::getFromOpaquePointer(function.function);
    contextBuilder.emplace(mlirFunc.getBody());
    context = new mlir::edsc::ScopedContext(*contextBuilder, mlirFunc.getLoc());
    return function;
  }

  void exit(py::object, py::object, py::object) {
    delete context;
    context = nullptr;
    contextBuilder.reset();
  }

  PythonFunction function;
  mlir::edsc::ScopedContext *context;
  llvm::Optional<OpBuilder> contextBuilder;
};

PythonFunctionContext PythonMLIRModule::makeNewFunctionContext(
    const std::string &name, const py::list &inputs,
    const std::vector<PythonType> &outputs, const py::kwargs &attributes) {
  auto func = declareFunction(name, inputs, outputs, attributes);
  func.define();
  return PythonFunctionContext(func);
}

PythonFunctionContext PythonMLIRModule::makeFunctionContext(
    PythonFunction f) {
  return PythonFunctionContext(f);
}

struct PythonBlockHandle {
  PythonBlockHandle() : value(nullptr) {}
  PythonBlockHandle(const PythonBlockHandle &other) = default;
  PythonBlockHandle(const mlir::edsc::BlockHandle &other) : value(other) {}
  operator mlir::edsc::BlockHandle() const { return value; }

  PythonType type(int index) { return types[index].getAsOpaquePointer(); }
  PythonValueHandle arg(int index) { return arguments[index]; }

  std::string str() {
    std::string s;
    llvm::raw_string_ostream os(s);
    value.getBlock()->print(os);
    return os.str();
  }

  mlir::edsc::BlockHandle value;
  std::vector<mlir::Type> types;
  std::vector<mlir::Value> arguments;
};

struct PythonOperationHandle {
  PythonOperationHandle() : value(nullptr) {}
  PythonOperationHandle(const PythonOperationHandle &other) = default;
  PythonOperationHandle(const mlir::edsc::OperationHandle &other)
      : value(other) {}
  operator OperationHandle() const { return value; }
  operator OperationHandle &() { return value; }
  PythonValueHandle getResult(int index) const {
    return value.getOperation()->getResult(index);
  }

  std::string str() const {
    std::string s;
    llvm::raw_string_ostream os(s);
    value.getOperation()->print(os);
    return os.str();
  }

  PythonOperationHandle call(const std::vector<PythonValueHandle> &args) {
    std::vector<Value> argValues;
    argValues.reserve(args.size());
    for (auto arg : args) argValues.push_back(arg.value);
    return OperationHandle::create<CallIndirectOp>(argValues[0], argValues);
  }

  mlir::edsc::OperationHandle value;
};

struct PythonLoopContext {
  PythonLoopContext(PythonValueHandle lb, PythonValueHandle ub, int64_t step)
      : lb(lb), ub(ub), step(step) {}
  PythonLoopContext(const PythonLoopContext &) = delete;
  PythonLoopContext(PythonLoopContext &&) = default;
  PythonLoopContext &operator=(const PythonLoopContext &) = delete;
  PythonLoopContext &operator=(PythonLoopContext &&) = default;
  ~PythonLoopContext() { assert(!builder && "did not exit from the context"); }

  PythonValueHandle enter() {
    mlir::Value iv;
    builder = new AffineLoopNestBuilder(&iv, lb.value, ub.value, step);
    return iv;
  }

  void exit(py::object, py::object, py::object) {
    (*builder)({});  // exit from the builder's scope.
    delete builder;
    builder = nullptr;
  }

  PythonValueHandle lb, ub;
  int64_t step;
  AffineLoopNestBuilder *builder = nullptr;
};

struct PythonLoopNestContext {
  PythonLoopNestContext(const std::vector<PythonValueHandle> &lbs,
                        const std::vector<PythonValueHandle> &ubs,
                        const std::vector<int64_t> steps)
      : lbs(lbs), ubs(ubs), steps(steps) {
    assert(lbs.size() == ubs.size() && lbs.size() == steps.size() &&
           "expected the same number of lower, upper bounds, and steps");
  }
  PythonLoopNestContext(const PythonLoopNestContext &) = delete;
  PythonLoopNestContext(PythonLoopNestContext &&) = default;
  PythonLoopNestContext &operator=(const PythonLoopNestContext &) = delete;
  PythonLoopNestContext &operator=(PythonLoopNestContext &&) = default;
  ~PythonLoopNestContext() {
    assert(!builder && "did not exit from the context");
  }

  std::vector<PythonValueHandle> enter() {
    if (steps.empty()) return {};

    std::vector<mlir::Value> handles(steps.size());
    builder = new AffineLoopNestBuilder(
        handles, std::vector<mlir::Value>(lbs.begin(), lbs.end()),
        std::vector<mlir::Value>(ubs.begin(), ubs.end()), steps);
    return std::vector<PythonValueHandle>(handles.begin(), handles.end());
  }

  void exit(py::object, py::object, py::object) {
    (*builder)({});  // exit from the builder's scope.
    delete builder;
    builder = nullptr;
  }

  std::vector<PythonValueHandle> lbs;
  std::vector<PythonValueHandle> ubs;
  std::vector<int64_t> steps;
  AffineLoopNestBuilder *builder = nullptr;
};

struct PythonBlockAppender {
  PythonBlockAppender(const PythonBlockHandle &handle) : handle(handle) {}
  PythonBlockHandle handle;
};

struct PythonBlockContext {
 public:
  PythonBlockContext() {
    createBlockBuilder();
    clearBuilder();
  }
  PythonBlockContext(const std::vector<PythonType> &argTypes) {
    handle.arguments.reserve(argTypes.size());
    for (const auto &t : argTypes) {
      auto type =
          Type::getFromOpaquePointer(reinterpret_cast<const void *>(t.type));
      handle.types.emplace_back(type);
      handle.arguments.push_back(Value());
    }
    createBlockBuilder();
    clearBuilder();
  }
  PythonBlockContext(const PythonBlockAppender &a) : handle(a.handle) {}
  PythonBlockContext(const PythonBlockContext &) = delete;
  PythonBlockContext(PythonBlockContext &&) = default;
  PythonBlockContext &operator=(const PythonBlockContext &) = delete;
  PythonBlockContext &operator=(PythonBlockContext &&) = default;
  ~PythonBlockContext() {
    assert(!builder && "did not exit from the block context");
  }

  // EDSC maintain an implicit stack of builders (mostly for keeping track of
  // insertion points); every operation gets inserted using the top-of-the-stack
  // builder.  Creating a new EDSC Builder automatically puts it on the stack,
  // effectively entering the block for it.
  void createBlockBuilder() {
    if (handle.value.getBlock())
      builder = new BlockBuilder(handle.value, mlir::edsc::Append());
    else
      builder = new BlockBuilder(&handle.value, handle.types, handle.arguments);
  }

  PythonBlockHandle enter() {
    createBlockBuilder();
    return handle;
  }

  void exit(py::object, py::object, py::object) { clearBuilder(); }

  PythonBlockHandle getHandle() { return handle; }

  // EDSC maintain an implicit stack of builders (mostly for keeping track of
  // insertion points); every operation gets inserted using the top-of-the-stack
  // builder.  Calling operator() on a builder pops the builder from the stack,
  // effectively resetting the insertion point to its position before we entered
  // the block.
  void clearBuilder() {
    (*builder)({});  // exit from the builder's scope.
    delete builder;
    builder = nullptr;
  }

  PythonBlockHandle handle;
  BlockBuilder *builder = nullptr;
};

struct PythonAttribute {
  PythonAttribute() : attr(nullptr) {}
  PythonAttribute(const mlir_attr_t &a) : attr(a) {}
  PythonAttribute(const PythonAttribute &other) = default;
  operator mlir_attr_t() { return attr; }

  operator Attribute() const { return Attribute::getFromOpaquePointer(attr); }

  std::string str() const {
    if (!attr) return "##null attr##";

    std::string res;
    llvm::raw_string_ostream os(res);
    Attribute().print(os);
    return res;
  }

  mlir_attr_t attr;
};

struct PythonAttributedType {
  PythonAttributedType() : type(nullptr) {}
  PythonAttributedType(mlir_type_t t) : type(t) {}
  PythonAttributedType(
      PythonType t,
      const std::unordered_map<std::string, PythonAttribute> &attributes =
          std::unordered_map<std::string, PythonAttribute>())
      : type(t), attrs(attributes) {}

  operator mlir_type_t() const { return type.type; }
  operator PythonType() const { return type; }

  // Return a vector of named attribute descriptors.  The vector owns the
  // mlir_named_attr_t objects it contains, but not the names and attributes
  // those objects point to (names and opaque pointers to attributes are owned
  // by `this`).
  std::vector<mlir_named_attr_t> getNamedAttrs() const {
    std::vector<mlir_named_attr_t> result;
    result.reserve(attrs.size());
    for (const auto &namedAttr : attrs)
      result.push_back({namedAttr.first.c_str(), namedAttr.second.attr});
    return result;
  }

  std::string str() {
    mlir::Type t = mlir::Type::getFromOpaquePointer(type);
    std::string res;
    llvm::raw_string_ostream os(res);
    t.print(os);
    if (attrs.empty()) return os.str();

    os << '{';
    bool first = true;
    for (const auto &namedAttr : attrs) {
      if (first)
        first = false;
      else
        os << ", ";
      os << namedAttr.first << ": " << namedAttr.second.str();
    }
    os << '}';

    return os.str();
  }

 private:
  PythonType type;
  std::unordered_map<std::string, PythonAttribute> attrs;
};

// Wraps mlir::AffineExpr.
struct PythonAffineExpr {
  PythonAffineExpr() : affine_expr() {}
  PythonAffineExpr(const AffineExpr &a) : affine_expr(a) {}
  PythonAffineExpr(const PythonAffineExpr &other) = default;

  operator AffineExpr() const { return affine_expr; }
  operator AffineExpr &() { return affine_expr; }

  AffineExpr get() const { return affine_expr; }

  std::string str() const {
    std::string res;
    llvm::raw_string_ostream os(res);
    affine_expr.print(os);
    return res;
  }

 private:
  AffineExpr affine_expr;
};

// Wraps mlir::AffineMap.
struct PythonAffineMap {
  PythonAffineMap() : affine_map() {}
  PythonAffineMap(const AffineMap &a) : affine_map(a) {}
  PythonAffineMap(const PythonAffineMap &other) = default;

  operator AffineMap() const { return affine_map; }
  operator AffineMap &() { return affine_map; }

  std::string str() const {
    std::string res;
    llvm::raw_string_ostream os(res);
    affine_map.print(os);
    return res;
  }

 private:
  AffineMap affine_map;
};

struct PythonIndexedValue {
  explicit PythonIndexedValue(const AffineIndexedValue &other)
      : indexed(other) {}
  PythonIndexedValue(PythonValueHandle handle) : indexed(handle.value) {}
  PythonIndexedValue(const PythonIndexedValue &other) = default;

  // Create a new indexed value with the same base as this one but with indices
  // provided as arguments.
  PythonIndexedValue index(const std::vector<PythonValueHandle> &indices) {
    std::vector<mlir::Value> handles(indices.begin(), indices.end());
    return PythonIndexedValue(AffineIndexedValue(indexed(handles)));
  }

  void store(const std::vector<PythonValueHandle> &indices,
             PythonValueHandle value) {
    // Uses the overloaded `operator=` to emit a store.
    index(indices).indexed = value.value;
  }

  PythonValueHandle load(const std::vector<PythonValueHandle> &indices) {
    // Uses the overloaded cast to `mlir::Value` to emit a load.
    return static_cast<mlir::Value>(index(indices).indexed);
  }

  AffineIndexedValue indexed;
};

template <typename ListTy, typename PythonTy, typename Ty>
ListTy makeCList(SmallVectorImpl<Ty> &owning, const py::list &list) {
  for (auto &inp : list) {
    owning.push_back(Ty{inp.cast<PythonTy>()});
  }
  return ListTy{owning.data(), owning.size()};
}

PythonFunction PythonMLIRModule::declareFunction(
    const std::string &name, const py::list &inputs,
    const std::vector<PythonType> &outputTypes,
    const py::kwargs &funcAttributes) {
  std::vector<PythonAttributedType> attributedInputs;
  attributedInputs.reserve(inputs.size());
  for (const auto &in : inputs) {
    std::string className = py::str(in.get_type());
    if (className.find(".Type'") != std::string::npos)
      attributedInputs.emplace_back(in.cast<PythonType>());
    else
      attributedInputs.push_back(in.cast<PythonAttributedType>());
  }

  // Create the function type.
  std::vector<mlir_type_t> ins(attributedInputs.begin(),
                               attributedInputs.end());
  std::vector<mlir_type_t> outs(outputTypes.begin(), outputTypes.end());
  auto funcType = ::makeFunctionType(
      mlir_context_t{&mlirContext}, mlir_type_list_t{ins.data(), ins.size()},
      mlir_type_list_t{outs.data(), outs.size()});

  // Build the list of function attributes.
  std::vector<mlir::NamedAttribute> attrs;
  attrs.reserve(funcAttributes.size());
  for (const auto &named : funcAttributes)
    attrs.emplace_back(
        Identifier::get(std::string(py::str(named.first)), &mlirContext),
        mlir::Attribute::getFromOpaquePointer(reinterpret_cast<const void *>(
            named.second.cast<PythonAttribute>().attr)));

  // Build the list of lists of function argument attributes.
  std::vector<mlir::NamedAttributeList> inputAttrs;
  inputAttrs.reserve(attributedInputs.size());
  for (const auto &in : attributedInputs) {
    std::vector<mlir::NamedAttribute> inAttrs;
    for (const auto &named : in.getNamedAttrs())
      inAttrs.emplace_back(Identifier::get(named.name, &mlirContext),
                           mlir::Attribute::getFromOpaquePointer(
                               reinterpret_cast<const void *>(named.value)));
    inputAttrs.emplace_back(inAttrs);
  }

  // Create the function itself.
  auto func = mlir::FuncOp::create(
      UnknownLoc::get(&mlirContext), name,
      mlir::Type::getFromOpaquePointer(funcType).cast<FunctionType>(), attrs,
      inputAttrs);
  symbolTable.insert(func);
  return func;
}

PythonAttributedType PythonType::attachAttributeDict(
    const std::unordered_map<std::string, PythonAttribute> &attrs) const {
  return PythonAttributedType(*this, attrs);
}

PythonAttribute PythonMLIRModule::integerAttr(PythonType type, int64_t value) {
  return PythonAttribute(::makeIntegerAttr(type, value));
}

PythonAttribute PythonMLIRModule::boolAttr(bool value) {
  return PythonAttribute(::makeBoolAttr(&mlirContext, value));
}

PythonAttribute PythonMLIRModule::floatAttr(float value) {
  return PythonAttribute(::makeFloatAttr(&mlirContext, value));
}

PythonAttribute PythonMLIRModule::stringAttr(const std::string &value) {
  return PythonAttribute(::makeStringAttr(&mlirContext, value.c_str()));
}

PythonAttribute PythonMLIRModule::arrayAttr(
    const std::vector<PythonAttribute> &values) {
  std::vector<mlir::Attribute> mlir_attributes(values.begin(), values.end());
  auto array_attr = ArrayAttr::get(
      llvm::ArrayRef<mlir::Attribute>(mlir_attributes), &mlirContext);
  return PythonAttribute(array_attr.getAsOpaquePointer());
}

PythonAttribute PythonMLIRModule::affineMapAttr(PythonAffineMap value) {
  return PythonAttribute(AffineMapAttr::get(value).getAsOpaquePointer());
}

PythonAttribute PythonMLIRModule::dictionaryAttr(
    const std::unordered_map<std::string, PythonAttribute> &attribute_map) {
  std::vector<NamedAttribute> attributes;
  attributes.reserve(attribute_map.size());
  for (const auto &pair : attribute_map) {
    attributes.emplace_back(Identifier::get(pair.first, &mlirContext),
                            Attribute::getFromOpaquePointer(pair.second.attr));
  }
  auto dictionaryAttribute = DictionaryAttr::get(attributes, &mlirContext);
  return PythonAttribute(dictionaryAttribute.getAsOpaquePointer());
}

PythonAffineExpr PythonMLIRModule::affineConstantExpr(int64_t value) {
  return PythonAffineExpr(getAffineConstantExpr(value, &mlirContext));
}

PythonAffineExpr PythonMLIRModule::affineSymbolExpr(unsigned position) {
  return PythonAffineExpr(getAffineSymbolExpr(position, &mlirContext));
}

PythonAffineExpr PythonMLIRModule::affineDimExpr(unsigned position) {
  return PythonAffineExpr(getAffineDimExpr(position, &mlirContext));
}

PythonAffineMap PythonMLIRModule::affineConstantMap(int64_t value) {
  return PythonAffineMap(AffineMap::getConstantMap(value, &mlirContext));
}

PythonAffineMap PythonMLIRModule::affineMap(
    unsigned dimCount, unsigned SymbolCount,
    const std::vector<PythonAffineExpr> &results) {
  std::vector<AffineExpr> mlir_results(results.begin(), results.end());
  return PythonAffineMap(
      AffineMap::get(dimCount, SymbolCount,
                     llvm::ArrayRef<AffineExpr>(mlir_results), &mlirContext));
}

void defineMlirEdscModule(py::module m) {
  m.doc() =
      "Python bindings for MLIR Embedded Domain-Specific Components (EDSCs)";
  m.def("version", []() { return "EDSC Python extensions v1.0"; });

  py::class_<PythonLoopContext>(
      m, "LoopContext", "A context for building the body of a 'for' loop")
      .def(py::init<PythonValueHandle, PythonValueHandle, int64_t>())
      .def("__enter__", &PythonLoopContext::enter)
      .def("__exit__", &PythonLoopContext::exit);

  py::class_<PythonLoopNestContext>(m, "LoopNestContext",
                                    "A context for building the body of a the "
                                    "innermost loop in a nest of 'for' loops")
      .def(py::init<const std::vector<PythonValueHandle> &,
                    const std::vector<PythonValueHandle> &,
                    const std::vector<int64_t> &>())
      .def("__enter__", &PythonLoopNestContext::enter)
      .def("__exit__", &PythonLoopNestContext::exit);

  m.def("constant_index", [](int64_t val) -> PythonValueHandle {
    return mlir::edsc::ValueBuilder<ConstantIndexOp>(val);
  });
  m.def("constant_int", [](int64_t val, int width) -> PythonValueHandle {
    return mlir::edsc::ValueBuilder<ConstantIntOp>(val, width);
  });
  m.def("constant_float", [](double val, PythonType type) -> PythonValueHandle {
    FloatType floatType =
        Type::getFromOpaquePointer(type.type).cast<FloatType>();
    assert(floatType);
    auto value = APFloat(val);
    bool lostPrecision;
    value.convert(floatType.getFloatSemantics(), APFloat::rmNearestTiesToEven,
                  &lostPrecision);
    return mlir::edsc::ValueBuilder<ConstantFloatOp>(value, floatType);
  });
  m.def("constant_function", [](PythonFunction func) -> PythonValueHandle {
    auto function = FuncOp::getFromOpaquePointer(func.function);
    auto attr = SymbolRefAttr::get(function.getName(), function.getContext());
    return mlir::edsc::ValueBuilder<ConstantOp>(function.getType(), attr);
  });
  m.def("appendTo", [](const PythonBlockHandle &handle) {
    return PythonBlockAppender(handle);
  });
  m.def(
      "ret",
      [](const std::vector<PythonValueHandle> &args) {
        std::vector<mlir::Value> values(args.begin(), args.end());
        (intrinsics::std_ret(ArrayRef<mlir::Value>{values}));  // vexing parse
        return PythonValueHandle(mlir::Value());
      },
      py::arg("args") = std::vector<PythonValueHandle>());
  m.def(
      "br",
      [](const PythonBlockHandle &dest,
         const std::vector<PythonValueHandle> &args) {
        std::vector<mlir::Value> values(args.begin(), args.end());
        intrinsics::std_br(dest, values);
        return PythonValueHandle(mlir::Value());
      },
      py::arg("dest"), py::arg("args") = std::vector<PythonValueHandle>());
  m.def(
      "cond_br",
      [](PythonValueHandle condition, const PythonBlockHandle &trueDest,
         const std::vector<PythonValueHandle> &trueArgs,
         const PythonBlockHandle &falseDest,
         const std::vector<PythonValueHandle> &falseArgs) -> PythonValueHandle {
        std::vector<mlir::Value> trueArguments(trueArgs.begin(),
                                               trueArgs.end());
        std::vector<mlir::Value> falseArguments(falseArgs.begin(),
                                                falseArgs.end());
        intrinsics::std_cond_br(condition, trueDest, trueArguments, falseDest,
                                falseArguments);
        return PythonValueHandle(mlir::Value());
      });
  m.def("index_cast",
        [](PythonValueHandle element, PythonType type) -> PythonValueHandle {
          return mlir::edsc::ValueBuilder<IndexCastOp>(
              element.value, Type::getFromOpaquePointer(type.type));
        });
  m.def("select",
        [](PythonValueHandle condition, PythonValueHandle trueValue,
           PythonValueHandle falseValue) -> PythonValueHandle {
          return mlir::edsc::ValueBuilder<SelectOp>(
              condition.value, trueValue.value, falseValue.value);
        });
  m.def("op",
        [](const std::string &name,
           const std::vector<PythonValueHandle> &operands,
           const std::vector<PythonType> &resultTypes,
           const py::kwargs &attributes) -> PythonValueHandle {
          std::vector<mlir::Value> operandHandles(operands.begin(),
                                                  operands.end());
          std::vector<Type> types;
          types.reserve(resultTypes.size());
          for (auto t : resultTypes)
            types.push_back(Type::getFromOpaquePointer(t.type));

          std::vector<NamedAttribute> attrs;
          attrs.reserve(attributes.size());
          for (const auto &a : attributes) {
            std::string name = py::str(a.first);
            auto pyAttr = a.second.cast<PythonAttribute>();
            auto cppAttr = Attribute::getFromOpaquePointer(pyAttr.attr);
            auto identifier =
                Identifier::get(name, ScopedContext::getContext());
            attrs.emplace_back(identifier, cppAttr);
          }
          return mlir::edsc::CustomValueBuilder(name, operandHandles, types,
                                                attrs);
        });
  m.def("multi_result_op",
        [](const std::string &name,
           const std::vector<PythonValueHandle> &operands,
           const std::vector<PythonType> &resultTypes,
           const py::kwargs &attributes) -> PythonOperationHandle {
          std::vector<mlir::Value> operandHandles(operands.begin(),
                                                  operands.end());
          std::vector<Type> types;
          types.reserve(resultTypes.size());
          for (auto t : resultTypes)
            types.push_back(Type::getFromOpaquePointer(t.type));

          std::vector<NamedAttribute> attrs;
          attrs.reserve(attributes.size());
          for (const auto &a : attributes) {
            std::string name = py::str(a.first);
            auto pyAttr = a.second.cast<PythonAttribute>();
            auto cppAttr = Attribute::getFromOpaquePointer(pyAttr.attr);
            auto identifier =
                Identifier::get(name, ScopedContext::getContext());
            attrs.emplace_back(identifier, cppAttr);
          }
          return OperationHandle::create(name, operandHandles, types, attrs);
        });

  py::class_<PythonFunction>(m, "Function", "Wrapping class for mlir::FuncOp.")
      .def(py::init<PythonFunction>())
      .def("__str__", &PythonFunction::str)
      .def("define", &PythonFunction::define,
           "Adds a body to the function if it does not already have one.  "
           "Returns true if the body was added")
      .def("arg", &PythonFunction::arg,
           "Get the mlir::Value to the indexed argument of the function");

  py::class_<PythonAttribute>(m, "Attribute",
                              "Wrapping class for mlir::Attribute")
      .def(py::init<PythonAttribute>())
      .def("__str__", &PythonAttribute::str);

  py::class_<PythonType>(m, "Type", "Wrapping class for mlir::Type.")
      .def(py::init<PythonType>())
      .def("__call__", &PythonType::attachAttributeDict,
           "Attach the attributes to these type, making it suitable for "
           "constructing functions with argument attributes")
      .def("__str__", &PythonType::str);

  py::class_<PythonAttributedType>(
      m, "AttributedType",
      "A class containing a wrapped mlir::Type and a wrapped "
      "mlir::NamedAttributeList that are used together, e.g. in function "
      "argument declaration")
      .def(py::init<PythonAttributedType>())
      .def("__str__", &PythonAttributedType::str);

  py::class_<PythonMLIRModule>(
      m, "MLIRModule",
      "An MLIRModule is the abstraction that owns the allocations to support "
      "compilation of a single mlir::ModuleOp into an ExecutionEngine backed "
      "by "
      "the LLVM ORC JIT. A typical flow consists in creating an MLIRModule, "
      "adding functions, compiling the module to obtain an ExecutionEngine on "
      "which named functions may be called. For now the only means to retrieve "
      "the ExecutionEngine is by calling `get_engine_address`. This mode of "
      "execution is limited to passing the pointer to C++ where the function "
      "is called. Extending the API to allow calling JIT compiled functions "
      "directly require integration with a tensor library (e.g. numpy). This "
      "is left as the prerogative of libraries and frameworks for now.")
      .def(py::init<>())
      .def("boolAttr", &PythonMLIRModule::boolAttr,
           "Creates an mlir::BoolAttr with the given value")
      .def(
          "integerAttr", &PythonMLIRModule::integerAttr,
          "Creates an mlir::IntegerAttr of the given type with the given value "
          "in the context associated with this MLIR module.")
      .def("floatAttr", &PythonMLIRModule::floatAttr,
           "Creates an mlir::FloatAttr with the given value")
      .def("stringAttr", &PythonMLIRModule::stringAttr,
           "Creates an mlir::StringAttr with the given value")
      .def("arrayAttr", &PythonMLIRModule::arrayAttr,
           "Creates an mlir::ArrayAttr of the given type with the given values "
           "in the context associated with this MLIR module.")
      .def("dictionaryAttr", &PythonMLIRModule::dictionaryAttr,
           "Creates a dictonary attribute with the given map.")
      .def("affineMapAttr", &PythonMLIRModule::affineMapAttr,
           "Creates an mlir::AffineMapAttr of the given type with the given "
           "value in the context associated with this MLIR module.")
      .def("declare_function", &PythonMLIRModule::declareFunction,
           "Declares a new mlir::FuncOp in the current mlir::ModuleOp.  The "
           "function arguments can have attributes.  The function has no "
           "definition and can be linked to an external library.")
      .def("make_function", &PythonMLIRModule::makeFunction,
           "Defines a new mlir::FuncOp in the current mlir::ModuleOp.")
      .def("new_function_context", &PythonMLIRModule::makeNewFunctionContext,
           "Defines a new mlir::FuncOp in the mlir::ModuleOp and creates the "
           "function context for building the body of the function.")
      .def("function_context", &PythonMLIRModule::makeFunctionContext,
           "Creates a function context for building the body of an existing "
           "function.")
      .def("get_function", &PythonMLIRModule::getNamedFunction,
           "Looks up the function with the given name in the module.")
      .def("make_memref_type", &PythonMLIRModule::makeMemRefType,
           "Returns an mlir::MemRefType of an elemental scalar. -1 is used to "
           "denote symbolic dimensions in the resulting memref shape.")
      .def("make_index_type", &PythonMLIRModule::makeIndexType,
           "Returns an mlir::IndexType")
      .def("make_type", &PythonMLIRModule::makeType,
           "Returns an mlir::Type defined by the IR passed in as the argument.")
      .def("get_ir", &PythonMLIRModule::getIR,
           "Returns a dump of the MLIR representation of the module. This is "
           "used for serde to support out-of-process execution as well as "
           "debugging purposes.")
      .def("affine_constant_expr", &PythonMLIRModule::affineConstantExpr,
           "Returns an affine constant expression.")
      .def("affine_symbol_expr", &PythonMLIRModule::affineSymbolExpr,
           "Returns an affine symbol expression.")
      .def("affine_dim_expr", &PythonMLIRModule::affineDimExpr,
           "Returns an affine dim expression.")
      .def("affine_constant_map", &PythonMLIRModule::affineConstantMap,
           "Returns an affine map with single constant result.")
      .def("affine_map", &PythonMLIRModule::affineMap, "Returns an affine map.",
           py::arg("dimCount"), py::arg("symbolCount"), py::arg("results"))
      .def("__str__", &PythonMLIRModule::getIR,
           "Get the string representation of the module");

  py::class_<PythonFunctionContext>(
      m, "FunctionContext", "A wrapper around mlir::edsc::ScopedContext")
      .def(py::init<PythonFunction>())
      .def("__enter__", &PythonFunctionContext::enter)
      .def("__exit__", &PythonFunctionContext::exit);

  {
    using namespace mlir::edsc::op;
    py::class_<PythonValueHandle>(m, "Value", "Wraps mlir::Value")
        .def(py::init<PythonValueHandle>())
        .def("__add__",
             [](PythonValueHandle lhs, PythonValueHandle rhs)
                 -> PythonValueHandle { return lhs.value + rhs.value; })
        .def("__sub__",
             [](PythonValueHandle lhs, PythonValueHandle rhs)
                 -> PythonValueHandle { return lhs.value - rhs.value; })
        .def("__mul__",
             [](PythonValueHandle lhs, PythonValueHandle rhs)
                 -> PythonValueHandle { return lhs.value * rhs.value; })
        .def("__div__",
             [](PythonValueHandle lhs, PythonValueHandle rhs)
                 -> PythonValueHandle { return lhs.value / rhs.value; })
        .def("__truediv__",
             [](PythonValueHandle lhs, PythonValueHandle rhs)
                 -> PythonValueHandle { return lhs.value / rhs.value; })
        .def("__floordiv__",
             [](PythonValueHandle lhs, PythonValueHandle rhs)
                 -> PythonValueHandle { return floorDiv(lhs, rhs); })
        .def("__mod__",
             [](PythonValueHandle lhs, PythonValueHandle rhs)
                 -> PythonValueHandle { return lhs.value % rhs.value; })
        .def("__lt__",
             [](PythonValueHandle lhs,
                PythonValueHandle rhs) -> PythonValueHandle {
               return mlir::edsc::ValueBuilder<CmpIOp>(CmpIPredicate::slt,
                                                       lhs.value, rhs.value);
             })
        .def("__le__",
             [](PythonValueHandle lhs,
                PythonValueHandle rhs) -> PythonValueHandle {
               return mlir::edsc::ValueBuilder<CmpIOp>(CmpIPredicate::sle,
                                                       lhs.value, rhs.value);
             })
        .def("__gt__",
             [](PythonValueHandle lhs,
                PythonValueHandle rhs) -> PythonValueHandle {
               return mlir::edsc::ValueBuilder<CmpIOp>(CmpIPredicate::sgt,
                                                       lhs.value, rhs.value);
             })
        .def("__ge__",
             [](PythonValueHandle lhs,
                PythonValueHandle rhs) -> PythonValueHandle {
               return mlir::edsc::ValueBuilder<CmpIOp>(CmpIPredicate::sge,
                                                       lhs.value, rhs.value);
             })
        .def("__eq__",
             [](PythonValueHandle lhs,
                PythonValueHandle rhs) -> PythonValueHandle {
               return mlir::edsc::ValueBuilder<CmpIOp>(CmpIPredicate::eq,
                                                       lhs.value, rhs.value);
             })
        .def("__ne__",
             [](PythonValueHandle lhs,
                PythonValueHandle rhs) -> PythonValueHandle {
               return mlir::edsc::ValueBuilder<CmpIOp>(CmpIPredicate::ne,
                                                       lhs.value, rhs.value);
             })
        .def("__invert__",
             [](PythonValueHandle handle) -> PythonValueHandle {
               return negate(handle.value);
             })
        .def("__and__",
             [](PythonValueHandle lhs, PythonValueHandle rhs)
                 -> PythonValueHandle { return lhs.value && rhs.value; })
        .def("__or__",
             [](PythonValueHandle lhs, PythonValueHandle rhs)
                 -> PythonValueHandle { return lhs.value || rhs.value; })
        .def("__call__", &PythonValueHandle::call)
        .def("type", &PythonValueHandle::type);
  }

  py::class_<PythonOperationHandle>(
      m, "OperationHandle", "A wrapper around mlir::edsc::OperationHandle")
      .def(py::init<PythonOperationHandle>())
      .def("__call__", &PythonOperationHandle::call)
      .def("result", &PythonOperationHandle::getResult);

  py::class_<PythonBlockAppender>(
      m, "BlockAppender",
      "A dummy class signaling BlockContext to append IR to the given "
      "block "
      "instead of creating a new block")
      .def(py::init<const PythonBlockHandle &>());
  py::class_<PythonBlockHandle>(m, "BlockHandle",
                                "A wrapper around mlir::edsc::BlockHandle")
      .def(py::init<PythonBlockHandle>())
      .def("arg", &PythonBlockHandle::arg);

  py::class_<PythonBlockContext>(m, "BlockContext",
                                 "A wrapper around mlir::edsc::BlockBuilder")
      .def(py::init<>())
      .def(py::init<const std::vector<PythonType> &>())
      .def(py::init<const PythonBlockAppender &>())
      .def("__enter__", &PythonBlockContext::enter)
      .def("__exit__", &PythonBlockContext::exit)
      .def("handle", &PythonBlockContext::getHandle);

  py::class_<PythonIndexedValue>(m, "IndexedValue",
                                 "A wrapper around mlir::edsc::IndexedValue")
      .def(py::init<PythonValueHandle>())
      .def("load", &PythonIndexedValue::load)
      .def("store", &PythonIndexedValue::store);

  py::class_<PythonAffineExpr>(m, "AffineExpr",
                               "A wrapper around mlir::AffineExpr")
      .def(py::init<PythonAffineExpr>())
      .def("__add__",
           [](PythonAffineExpr lhs, int64_t rhs) -> PythonAffineExpr {
             return PythonAffineExpr(lhs.get() + rhs);
           })
      .def("__add__",
           [](PythonAffineExpr lhs, PythonAffineExpr rhs) -> PythonAffineExpr {
             return PythonAffineExpr(lhs.get() + rhs.get());
           })
      .def("__neg__",
           [](PythonAffineExpr lhs) -> PythonAffineExpr {
             return PythonAffineExpr(-lhs.get());
           })
      .def("__sub__",
           [](PythonAffineExpr lhs, int64_t rhs) -> PythonAffineExpr {
             return PythonAffineExpr(lhs.get() - rhs);
           })
      .def("__sub__",
           [](PythonAffineExpr lhs, PythonAffineExpr rhs) -> PythonAffineExpr {
             return PythonAffineExpr(lhs.get() - rhs.get());
           })
      .def("__mul__",
           [](PythonAffineExpr lhs, int64_t rhs) -> PythonAffineExpr {
             return PythonAffineExpr(lhs.get() * rhs);
           })
      .def("__mul__",
           [](PythonAffineExpr lhs, PythonAffineExpr rhs) -> PythonAffineExpr {
             return PythonAffineExpr(lhs.get() * rhs.get());
           })
      .def("__floordiv__",
           [](PythonAffineExpr lhs, uint64_t rhs) -> PythonAffineExpr {
             return PythonAffineExpr(lhs.get().floorDiv(rhs));
           })
      .def("__floordiv__",
           [](PythonAffineExpr lhs, PythonAffineExpr rhs) -> PythonAffineExpr {
             return PythonAffineExpr(lhs.get().floorDiv(rhs.get()));
           })
      .def("ceildiv",
           [](PythonAffineExpr lhs, uint64_t rhs) -> PythonAffineExpr {
             return PythonAffineExpr(lhs.get().ceilDiv(rhs));
           })
      .def("ceildiv",
           [](PythonAffineExpr lhs, PythonAffineExpr rhs) -> PythonAffineExpr {
             return PythonAffineExpr(lhs.get().ceilDiv(rhs.get()));
           })
      .def("__mod__",
           [](PythonAffineExpr lhs, uint64_t rhs) -> PythonAffineExpr {
             return PythonAffineExpr(lhs.get() % rhs);
           })
      .def("__mod__",
           [](PythonAffineExpr lhs, PythonAffineExpr rhs) -> PythonAffineExpr {
             return PythonAffineExpr(lhs.get() % rhs.get());
           })
      .def("compose",
           [](PythonAffineExpr self, PythonAffineMap map) -> PythonAffineExpr {
             return PythonAffineExpr(self.get().compose(map));
           })
      .def(
          "get_constant_value",
          [](PythonAffineExpr self) -> py::object {
            auto const_expr = self.get().dyn_cast<AffineConstantExpr>();
            if (const_expr) return py::cast(const_expr.getValue());
            return py::none();
          },
          "Returns the constant value for the affine expression if any, or "
          "returns None.")
      .def("__str__", &PythonAffineExpr::str);

  py::class_<PythonAffineMap>(m, "AffineMap",
                              "A wrapper around mlir::AffineMap")
      .def(py::init<PythonAffineMap>())
      .def("__str__", &PythonAffineMap::str);
}

} // namespace edsc
} // namespace mlir