mlir/lib/Bindings/Python/DialectTransform.cpp

//===- DialectTransform.cpp - 'transform' dialect submodule ---------------===//
//
// 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 <string>

#include "Rewrite.h"
#include "mlir-c/Dialect/Transform.h"
#include "mlir-c/IR.h"
#include "mlir-c/Rewrite.h"
#include "mlir-c/Support.h"
#include "mlir/Bindings/Python/IRCore.h"
#include "mlir/Bindings/Python/IRInterfaces.h"
#include "nanobind/nanobind.h"
#include <nanobind/trampoline.h>

namespace nb = nanobind;
using namespace mlir::python::nanobind_adaptors;

namespace mlir {
namespace python {
namespace MLIR_BINDINGS_PYTHON_DOMAIN {
namespace transform {

//===----------------------------------------------------------------------===//
// TransformRewriter
//===----------------------------------------------------------------------===//
class PyTransformRewriter : public PyRewriterBase<PyTransformRewriter> {
public:
  static constexpr const char *pyClassName = "TransformRewriter";

  PyTransformRewriter(MlirTransformRewriter rewriter)
      : PyRewriterBase(mlirTransformRewriterAsBase(rewriter)) {}
};

//===----------------------------------------------------------------------===//
// TransformResults
//===----------------------------------------------------------------------===//
class PyTransformResults {
public:
  PyTransformResults(MlirTransformResults results) : results(results) {}

  MlirTransformResults get() const { return results; }

  void setOps(PyValue &result,
              const nb::typed<nb::sequence, PyOperationBase> &ops) {
    std::vector<MlirOperation> opsVec;
    opsVec.reserve(nb::len(ops));
    for (auto op : ops) {
      opsVec.push_back(nb::cast<MlirOperation>(op));
    }
    mlirTransformResultsSetOps(results, result, opsVec.size(), opsVec.data());
  }

  void setValues(PyValue &result,
                 const nb::typed<nb::sequence, PyValue> &values) {
    std::vector<MlirValue> valuesVec;
    valuesVec.reserve(nb::len(values));
    for (auto item : values) {
      valuesVec.push_back(nb::cast<MlirValue>(item));
    }
    mlirTransformResultsSetValues(results, result, valuesVec.size(),
                                  valuesVec.data());
  }

  void setParams(PyValue &result,
                 const nb::typed<nb::sequence, PyAttribute> &params) {
    std::vector<MlirAttribute> paramsVec;
    paramsVec.reserve(nb::len(params));
    for (auto item : params) {
      paramsVec.push_back(nb::cast<MlirAttribute>(item));
    }
    mlirTransformResultsSetParams(results, result, paramsVec.size(),
                                  paramsVec.data());
  }

  static void bind(nanobind::module_ &m) {
    nb::class_<PyTransformResults>(m, "TransformResults")
        .def(nb::init<MlirTransformResults>())
        .def("set_ops", &PyTransformResults::setOps,
             "Set the payload operations for a transform result.",
             nb::arg("result"), nb::arg("ops"))
        .def("set_values", &PyTransformResults::setValues,
             "Set the payload values for a transform result.",
             nb::arg("result"), nb::arg("values"))
        .def("set_params", &PyTransformResults::setParams,
             "Set the parameters for a transform result.", nb::arg("result"),
             nb::arg("params"));
  }

private:
  MlirTransformResults results;
};

//===----------------------------------------------------------------------===//
// TransformState
//===----------------------------------------------------------------------===//
class PyTransformState {
public:
  PyTransformState(MlirTransformState state) : state(state) {}

  MlirTransformState get() const { return state; }

  static void bind(nanobind::module_ &m) {
    nb::class_<PyTransformState>(m, "TransformState")
        .def(nb::init<MlirTransformState>())
        .def("get_payload_ops", &PyTransformState::getPayloadOps,
             "Get the payload operations associated with a transform IR value.",
             nb::arg("operand"))
        .def("get_payload_values", &PyTransformState::getPayloadValues,
             "Get the payload values associated with a transform IR value.",
             nb::arg("operand"))
        .def("get_params", &PyTransformState::getParams,
             "Get the parameters (attributes) associated with a transform IR "
             "value.",
             nb::arg("operand"));
  }

private:
  nanobind::list getPayloadOps(PyValue &value) {
    nanobind::list result;
    mlirTransformStateForEachPayloadOp(
        state, value,
        [](MlirOperation op, void *userData) {
          PyMlirContextRef context =
              PyMlirContext::forContext(mlirOperationGetContext(op));
          auto opview = PyOperation::forOperation(context, op)->createOpView();
          static_cast<nanobind::list *>(userData)->append(opview);
        },
        &result);
    return result;
  }

  nanobind::list getPayloadValues(PyValue &value) {
    nanobind::list result;
    mlirTransformStateForEachPayloadValue(
        state, value,
        [](MlirValue val, void *userData) {
          static_cast<nanobind::list *>(userData)->append(val);
        },
        &result);
    return result;
  }

  nanobind::list getParams(PyValue &value) {
    nanobind::list result;
    mlirTransformStateForEachParam(
        state, value,
        [](MlirAttribute attr, void *userData) {
          static_cast<nanobind::list *>(userData)->append(attr);
        },
        &result);
    return result;
  }

  MlirTransformState state;
};

//===----------------------------------------------------------------------===//
// TransformOpInterface
//===----------------------------------------------------------------------===//
class PyTransformOpInterface
    : public PyConcreteOpInterface<PyTransformOpInterface> {
public:
  using PyConcreteOpInterface<PyTransformOpInterface>::PyConcreteOpInterface;

  constexpr static const char *pyClassName = "TransformOpInterface";
  constexpr static GetTypeIDFunctionTy getInterfaceID =
      &mlirTransformOpInterfaceTypeID;

  /// Attach a new TransformOpInterface FallbackModel to the named operation.
  /// The FallbackModel acts as a trampoline for callbacks on the Python class.
  static void attach(nb::object &target, const std::string &opName,
                     DefaultingPyMlirContext ctx) {
    // Prepare the callbacks that will be used by the FallbackModel.
    MlirTransformOpInterfaceCallbacks callbacks;
    // Make the pointer to the Python class available to the callbacks.
    callbacks.userData = target.ptr();
    nb::handle(static_cast<PyObject *>(callbacks.userData)).inc_ref();

    // The above ref bump is all we need as initialization, no need to run the
    // construct callback.
    callbacks.construct = nullptr;
    // Upon the FallbackModel's destruction, drop the ref to the Python class.
    callbacks.destruct = [](void *userData) {
      nb::handle(static_cast<PyObject *>(userData)).dec_ref();
    };
    // The apply callback which calls into Python.
    callbacks.apply = [](MlirOperation op, MlirTransformRewriter rewriter,
                         MlirTransformResults results, MlirTransformState state,
                         void *userData) -> MlirDiagnosedSilenceableFailure {
      nb::handle pyClass(static_cast<PyObject *>(userData));

      auto pyApply = nb::cast<nb::callable>(nb::getattr(pyClass, "apply"));

      auto pyRewriter = PyTransformRewriter(rewriter);
      auto pyResults = PyTransformResults(results);
      auto pyState = PyTransformState(state);

      // Invoke `pyClass.apply(opview(op), rewriter, results, state)` as a
      // staticmethod.
      PyMlirContextRef context =
          PyMlirContext::forContext(mlirOperationGetContext(op));
      auto opview = PyOperation::forOperation(context, op)->createOpView();
      nb::object res = pyApply(opview, pyRewriter, pyResults, pyState);

      return nb::cast<MlirDiagnosedSilenceableFailure>(res);
    };

    // The allows_repeated_handle_operands callback which calls into Python.
    callbacks.allowsRepeatedHandleOperands = [](MlirOperation op,
                                                void *userData) -> bool {
      nb::handle pyClass(static_cast<PyObject *>(userData));

      auto pyAllowRepeatedHandleOperands = nb::cast<nb::callable>(
          nb::getattr(pyClass, "allow_repeated_handle_operands"));

      // Invoke `pyClass.allow_repeated_handle_operands(opview(op))` as a
      // staticmethod.
      PyMlirContextRef context =
          PyMlirContext::forContext(mlirOperationGetContext(op));
      auto opview = PyOperation::forOperation(context, op)->createOpView();
      nb::object res = pyAllowRepeatedHandleOperands(opview);

      return nb::cast<bool>(res);
    };

    // Attach a FallbackModel, which calls into Python, to the named operation.
    mlirTransformOpInterfaceAttachFallbackModel(
        ctx->get(), mlirStringRefCreate(opName.c_str(), opName.size()),
        callbacks);
  }

  static void bindDerived(ClassTy &cls) {
    cls.attr("attach") = classmethod(
        [](const nb::object &cls, const nb::object &opName, nb::object target,
           DefaultingPyMlirContext context) {
          if (target.is_none())
            target = cls;
          return attach(target, nb::cast<std::string>(opName), context);
        },
        nb::arg("cls"), nb::arg("op_name"), nb::kw_only(),
        nb::arg("target").none() = nb::none(),
        nb::arg("context").none() = nb::none(),
        "Attach the interface subclass to the given operation name.");
  }
};

//===----------------------------------------------------------------------===//
// PatternDescriptorOpInterface
//===----------------------------------------------------------------------===//
class PyPatternDescriptorOpInterface
    : public PyConcreteOpInterface<PyPatternDescriptorOpInterface> {
public:
  using PyConcreteOpInterface<
      PyPatternDescriptorOpInterface>::PyConcreteOpInterface;

  constexpr static const char *pyClassName = "PatternDescriptorOpInterface";
  constexpr static GetTypeIDFunctionTy getInterfaceID =
      &mlirPatternDescriptorOpInterfaceTypeID;

  /// Attach a new PatternDescriptorOpInterface FallbackModel to the named
  /// operation. The FallbackModel acts as a trampoline for callbacks on the
  /// Python class.
  static void attach(nb::object &target, const std::string &opName,
                     DefaultingPyMlirContext ctx) {
    // Prepare the callbacks that will be used by the FallbackModel.
    MlirPatternDescriptorOpInterfaceCallbacks callbacks;
    // Make the pointer to the Python class available to the callbacks.
    callbacks.userData = target.ptr();
    nb::handle(static_cast<PyObject *>(callbacks.userData)).inc_ref();

    // The above ref bump is all we need as initialization, no need to run the
    // construct callback.
    callbacks.construct = nullptr;
    // Upon the FallbackModel's destruction, drop the ref to the Python class.
    callbacks.destruct = [](void *userData) {
      nb::handle(static_cast<PyObject *>(userData)).dec_ref();
    };

    // The populatePatterns callback which calls into Python.
    callbacks.populatePatterns =
        [](MlirOperation op, MlirRewritePatternSet patterns, void *userData) {
          nb::handle pyClass(static_cast<PyObject *>(userData));

          auto pyPopulatePatterns =
              nb::cast<nb::callable>(nb::getattr(pyClass, "populate_patterns"));

          auto pyPatterns = PyRewritePatternSet(patterns);

          // Invoke `pyClass.populate_patterns(opview(op), patterns)` as a
          // staticmethod.
          MlirContext ctx = mlirOperationGetContext(op);
          PyMlirContextRef context = PyMlirContext::forContext(ctx);
          auto opview = PyOperation::forOperation(context, op)->createOpView();
          pyPopulatePatterns(opview, pyPatterns);
        };

    // The populatePatternsWithState callback which calls into Python.
    // Check if the Python class has populate_patterns_with_state method.
    if (nb::hasattr(target, "populate_patterns_with_state")) {
      callbacks.populatePatternsWithState = [](MlirOperation op,
                                               MlirRewritePatternSet patterns,
                                               MlirTransformState state,
                                               void *userData) {
        nb::handle pyClass(static_cast<PyObject *>(userData));

        auto pyPopulatePatternsWithState = nb::cast<nb::callable>(
            nb::getattr(pyClass, "populate_patterns_with_state"));

        auto pyPatterns = PyRewritePatternSet(patterns);
        auto pyState = PyTransformState(state);

        // Invoke `pyClass.populate_patterns_with_state(opview(op), patterns,
        // state)` as a staticmethod.
        MlirContext ctx = mlirOperationGetContext(op);
        PyMlirContextRef context = PyMlirContext::forContext(ctx);
        auto opview = PyOperation::forOperation(context, op)->createOpView();
        pyPopulatePatternsWithState(opview, pyPatterns, pyState);
      };
    } else {
      // Use default implementation (will call populatePatterns).
      callbacks.populatePatternsWithState = nullptr;
    }

    // Attach a FallbackModel, which calls into Python, to the named operation.
    mlirPatternDescriptorOpInterfaceAttachFallbackModel(
        ctx->get(), mlirStringRefCreate(opName.c_str(), opName.size()),
        callbacks);
  }

  static void bindDerived(ClassTy &cls) {
    cls.attr("attach") = classmethod(
        [](const nb::object &cls, const nb::object &opName, nb::object target,
           DefaultingPyMlirContext context) {
          if (target.is_none())
            target = cls;
          return attach(target, nb::cast<std::string>(opName), context);
        },
        nb::arg("cls"), nb::arg("op_name"), nb::kw_only(),
        nb::arg("target").none() = nb::none(),
        nb::arg("context").none() = nb::none(),
        "Attach the interface subclass to the given operation name.");
  }
};

//===-------------------------------------------------------------------===//
// AnyOpType
//===-------------------------------------------------------------------===//

struct AnyOpType : PyConcreteType<AnyOpType> {
  static constexpr IsAFunctionTy isaFunction = mlirTypeIsATransformAnyOpType;
  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
      mlirTransformAnyOpTypeGetTypeID;
  static constexpr const char *pyClassName = "AnyOpType";
  static inline const MlirStringRef name = mlirTransformAnyOpTypeGetName();
  using Base::Base;

  static void bindDerived(ClassTy &c) {
    c.def_static(
        "get",
        [](DefaultingPyMlirContext context) {
          return AnyOpType(context->getRef(),
                           mlirTransformAnyOpTypeGet(context.get()->get()));
        },
        "Get an instance of AnyOpType in the given context.",
        nb::arg("context").none() = nb::none());
  }
};

//===-------------------------------------------------------------------===//
// AnyParamType
//===-------------------------------------------------------------------===//

struct AnyParamType : PyConcreteType<AnyParamType> {
  static constexpr IsAFunctionTy isaFunction = mlirTypeIsATransformAnyParamType;
  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
      mlirTransformAnyParamTypeGetTypeID;
  static constexpr const char *pyClassName = "AnyParamType";
  static inline const MlirStringRef name = mlirTransformAnyParamTypeGetName();
  using Base::Base;

  static void bindDerived(ClassTy &c) {
    c.def_static(
        "get",
        [](DefaultingPyMlirContext context) {
          return AnyParamType(context->getRef(), mlirTransformAnyParamTypeGet(
                                                     context.get()->get()));
        },
        "Get an instance of AnyParamType in the given context.",
        nb::arg("context").none() = nb::none());
  }
};

//===-------------------------------------------------------------------===//
// AnyValueType
//===-------------------------------------------------------------------===//

struct AnyValueType : PyConcreteType<AnyValueType> {
  static constexpr IsAFunctionTy isaFunction = mlirTypeIsATransformAnyValueType;
  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
      mlirTransformAnyValueTypeGetTypeID;
  static constexpr const char *pyClassName = "AnyValueType";
  static inline const MlirStringRef name = mlirTransformAnyValueTypeGetName();
  using Base::Base;

  static void bindDerived(ClassTy &c) {
    c.def_static(
        "get",
        [](DefaultingPyMlirContext context) {
          return AnyValueType(context->getRef(), mlirTransformAnyValueTypeGet(
                                                     context.get()->get()));
        },
        "Get an instance of AnyValueType in the given context.",
        nb::arg("context").none() = nb::none());
  }
};

//===-------------------------------------------------------------------===//
// OperationType
//===-------------------------------------------------------------------===//

struct OperationType : PyConcreteType<OperationType> {
  static constexpr IsAFunctionTy isaFunction =
      mlirTypeIsATransformOperationType;
  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
      mlirTransformOperationTypeGetTypeID;
  static constexpr const char *pyClassName = "OperationType";
  static inline const MlirStringRef name = mlirTransformOperationTypeGetName();
  using Base::Base;

  static void bindDerived(ClassTy &c) {
    c.def_static(
        "get",
        [](const std::string &operationName, DefaultingPyMlirContext context) {
          MlirStringRef cOperationName =
              mlirStringRefCreate(operationName.data(), operationName.size());
          return OperationType(context->getRef(),
                               mlirTransformOperationTypeGet(
                                   context.get()->get(), cOperationName));
        },
        "Get an instance of OperationType for the given kind in the given "
        "context",
        nb::arg("operation_name"), nb::arg("context").none() = nb::none());
    c.def_prop_ro(
        "operation_name",
        [](const OperationType &type) {
          MlirStringRef operationName =
              mlirTransformOperationTypeGetOperationName(type);
          return nb::str(operationName.data, operationName.length);
        },
        "Get the name of the payload operation accepted by the handle.");
  }
};

//===-------------------------------------------------------------------===//
// ParamType
//===-------------------------------------------------------------------===//

struct ParamType : PyConcreteType<ParamType> {
  static constexpr IsAFunctionTy isaFunction = mlirTypeIsATransformParamType;
  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
      mlirTransformParamTypeGetTypeID;
  static constexpr const char *pyClassName = "ParamType";
  static inline const MlirStringRef name = mlirTransformParamTypeGetName();
  using Base::Base;

  static void bindDerived(ClassTy &c) {
    c.def_static(
        "get",
        [](const PyType &type, DefaultingPyMlirContext context) {
          return ParamType(context->getRef(), mlirTransformParamTypeGet(
                                                  context.get()->get(), type));
        },
        "Get an instance of ParamType for the given type in the given context.",
        nb::arg("type"), nb::arg("context").none() = nb::none());
    c.def_prop_ro(
        "type",
        [](ParamType type) {
          return PyType(type.getContext(), mlirTransformParamTypeGetType(type))
              .maybeDownCast();
        },
        "Get the type this ParamType is associated with.");
  }
};

//===----------------------------------------------------------------------===//
// MemoryEffectsOpInterface helpers
//===----------------------------------------------------------------------===//

namespace {
void onlyReadsHandle(nb::iterable &operands,
                     PyMemoryEffectsInstanceList effects) {
  std::vector<MlirOpOperand> operandsVec;
  for (auto operand : operands)
    operandsVec.push_back(nb::cast<PyOpOperand>(operand));
  mlirTransformOnlyReadsHandle(operandsVec.data(), operandsVec.size(),
                               effects.effects);
};

void consumesHandle(nb::iterable &operands,
                    PyMemoryEffectsInstanceList effects) {
  std::vector<MlirOpOperand> operandsVec;
  for (auto operand : operands)
    operandsVec.push_back(nb::cast<PyOpOperand>(operand));
  mlirTransformConsumesHandle(operandsVec.data(), operandsVec.size(),
                              effects.effects);
};

void producesHandle(nb::iterable &results,
                    PyMemoryEffectsInstanceList effects) {
  std::vector<MlirValue> resultsVec;
  for (auto result : results)
    resultsVec.push_back(nb::cast<PyOpResult>(result).get());
  mlirTransformProducesHandle(resultsVec.data(), resultsVec.size(),
                              effects.effects);
};

void modifiesPayload(PyMemoryEffectsInstanceList effects) {
  mlirTransformModifiesPayload(effects.effects);
}

void onlyReadsPayload(PyMemoryEffectsInstanceList effects) {
  mlirTransformOnlyReadsPayload(effects.effects);
}
} // namespace

static void populateDialectTransformSubmodule(nb::module_ &m) {
  nb::enum_<MlirDiagnosedSilenceableFailure>(m, "DiagnosedSilenceableFailure")
      .value("Success", MlirDiagnosedSilenceableFailureSuccess)
      .value("SilenceableFailure",
             MlirDiagnosedSilenceableFailureSilenceableFailure)
      .value("DefiniteFailure", MlirDiagnosedSilenceableFailureDefiniteFailure);

  AnyOpType::bind(m);
  AnyParamType::bind(m);
  AnyValueType::bind(m);
  OperationType::bind(m);
  ParamType::bind(m);

  PyTransformRewriter::bind(m);
  PyTransformResults::bind(m);
  PyTransformState::bind(m);
  PyTransformOpInterface::bind(m);
  PyPatternDescriptorOpInterface::bind(m);

  m.def("only_reads_handle", onlyReadsHandle,
        "Mark operands as only reading handles.", nb::arg("operands"),
        nb::arg("effects"));

  m.def("consumes_handle", consumesHandle,
        "Mark operands as consuming handles.", nb::arg("operands"),
        nb::arg("effects"));

  m.def("produces_handle", producesHandle, "Mark results as producing handles.",
        nb::arg("results"), nb::arg("effects"));

  m.def("modifies_payload", modifiesPayload,
        "Mark the transform as modifying the payload.", nb::arg("effects"));

  m.def("only_reads_payload", onlyReadsPayload,
        "Mark the transform as only reading the payload.", nb::arg("effects"));
}
} // namespace transform
} // namespace MLIR_BINDINGS_PYTHON_DOMAIN
} // namespace python
} // namespace mlir

NB_MODULE(_mlirDialectsTransform, m) {
  m.doc() = "MLIR Transform dialect.";
  mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::transform::
      populateDialectTransformSubmodule(m);
}