Skip to content

Commit

Permalink
Environment variables to control MLIR execution (openvinotoolkit#154)
Browse files Browse the repository at this point in the history 
* Check MatMul expected attributes in a predicate instead of assert in the transformation callback. Fixes PyTorch addmm layer tests.

* Put rank == 2 restriction on MatMul conversion. Fixes PyTorch linear layer tests.

* Enable MLIR and TPP-MLIR activation via environment variables

* Added debug macros controlled by OV_MLIR_DEBUG env variable. Disabled debug prints by default.
  • Loading branch information
@slyalin
slyalin authored Jul 26, 2024
1 parent 7fba0a4 commit c793696
Show file tree
Hide file tree
Showing 5 changed files with 155 additions and 109 deletions.
80 changes: 57 additions & 23 deletions src/common/transformations/src/transformations/mlir/convert.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
#include <openvino/pass/graph_rewrite.hpp>
#include <openvino/pass/manager.hpp>
#include <openvino/pass/pattern/op/wrap_type.hpp>
#include <openvino/util/env_util.hpp>
#include <unordered_map>

// TODO: Prune unused headers -- it's hard to understand needed ones
Expand Down Expand Up @@ -187,7 +188,7 @@ mlir::OwningOpRef<mlir::ModuleOp> ngraph_to_mlir(MLIRContext* context,


// This pass converts a group of nodes into a single MLIROp
NodePtr ngraph_to_mlir_op(MLIRContext* context, SubgraphPtr subgraph) {
NodePtr ngraph_to_mlir_op(MLIRContext* context, SubgraphPtr subgraph, bool tpp_mlir_enabled) {
mlir::OwningOpRef<mlir::ModuleOp> module = ngraph_to_mlir(context, subgraph->inputs, subgraph->nodes, subgraph->outputs);

const auto& inputs = subgraph->inputs;
Expand All @@ -203,8 +204,9 @@ NodePtr ngraph_to_mlir_op(MLIRContext* context, SubgraphPtr subgraph) {
if(0 == input_map.count(symbol)) {
input_map[symbol] = Index(i, j);
} else {
std::cerr << "[ DEBUG ] Lost equality constraint for dimensions in output " << input << "\n"
<< " If the constraint is violated in runtime it will result in the undefined behaviour.\n";
OPENVINO_MLIR_DEBUG_PRINT(
"[ DEBUG ] Lost equality constraint for dimensions in output " << input << ".\n" <<
" If the constraint is violated in runtime it will result in the undefined behaviour.\n");
}
}
}
Expand All @@ -230,7 +232,7 @@ NodePtr ngraph_to_mlir_op(MLIRContext* context, SubgraphPtr subgraph) {
}
return std::make_shared<MLIROp>(
subgraph->inputs,
std::make_shared<MLIREvaluate>(std::move(module)),
std::make_shared<MLIREvaluate>(std::move(module), tpp_mlir_enabled),
output_types,
output_map
);
Expand All @@ -251,16 +253,20 @@ void replace_subgraph(SubgraphPtr subgraph, NodePtr node) {

class Partitioner : public ov::pass::ModelPass {
MLIRContext* context;
bool tpp_mlir_enabled;
public:
OPENVINO_RTTI("Partitioner");

Partitioner(MLIRContext* context) : context(context) {}
Partitioner(MLIRContext* context, bool tpp_mlir_enabled) :
context(context),
tpp_mlir_enabled(tpp_mlir_enabled)
{}

bool run_on_model(const std::shared_ptr<ov::Model>& model) override {
SubgraphTracker tracker([this](SubgraphPtr subgraph) {
auto mlir_op = ngraph_to_mlir_op(context, subgraph);
auto mlir_op = ngraph_to_mlir_op(context, subgraph, tpp_mlir_enabled);
replace_subgraph(subgraph, mlir_op);
std::cerr << "Created MLIR op: " << mlir_op << "\n";
OPENVINO_MLIR_DEBUG_PRINT("Created MLIR op: " << mlir_op << "\n");
}
);
for(auto node: model->get_ordered_ops()) {
Expand All @@ -272,7 +278,7 @@ class Partitioner : public ov::pass::ModelPass {
};


void injectMLIR(std::shared_ptr<ov::Model> model, MLIRContext* context) {
void injectMLIR(std::shared_ptr<ov::Model> model, MLIRContext* context, bool tpp_mlir_enabled) {
ov::pass::Manager manager;
using namespace ov::op;
manager.set_per_pass_validation(false);
Expand All @@ -283,42 +289,55 @@ void injectMLIR(std::shared_ptr<ov::Model> model, MLIRContext* context) {
manager.register_pass<BinaryEltwisePattern<v1::Divide, linalg::DivOp>>(ov::element::f32);
manager.register_pass<ReluPattern>();
manager.register_pass<MatMulPattern>();
manager.register_pass<Partitioner>(context);
manager.register_pass<Partitioner>(context, tpp_mlir_enabled);
manager.run_passes(model);
model->validate_nodes_and_infer_types();
}


MLIRContext* get_shared_mlir_context() {
MLIRContext* get_shared_mlir_context(bool tpp_mlir_enabled_current) {
// Gives MLIRContext instance shared for entire OV process and initialized once upon the initial request
// FIXME: Bind with OpenVINO lifetime in the sutable class instead of dirty tricking with static lifetime

static std::shared_ptr<MLIRContext> context;
static bool tpp_mlir_enabled = tpp_mlir_enabled_current;

if(context) {
if(tpp_mlir_enabled_current != tpp_mlir_enabled) {
OPENVINO_MLIR_DEBUG_PRINT("[ DEBUG ] Switched TPP mode, reinitialize MLIR context\n");
tpp_mlir_enabled = tpp_mlir_enabled_current;
context.reset();
}
}

if (!context) {

// Initialize the LLVM machinery
llvm::InitializeNativeTarget();
llvm::InitializeNativeTargetAsmPrinter();

std::cerr << "[ DEBUG ] Using TPP_MLIR: ";
#if TPP_MLIR
OPENVINO_MLIR_DEBUG_PRINT("[ DEBUG ] Using TPP_MLIR: ");
if(tpp_mlir_enabled) {
OPENVINO_MLIR_DEBUG_PRINT("YES\n");
// Initialize GPU-related LLVM machinery
tpp::initializeGpuTargets();
std::cerr << "YES\n";
#else
std::cerr << "NO\n";
#endif
#ifdef TPP_MLIR
tpp::initializeGpuTargets();
#endif
} else {
OPENVINO_MLIR_DEBUG_PRINT("NO\n");
}

// Add the following to include *all* MLIR Core dialects, or selectively
// include what you need like above. You only need to register dialects that
// will be *parsed* by the tool, not the one generated
DialectRegistry registry;
#if TPP_MLIR
registry.insert<mlir::xsmm::XsmmDialect>();
registry.insert<mlir::check::CheckDialect>();
registry.insert<mlir::perf::PerfDialect>();
#endif
if(tpp_mlir_enabled) {
#ifdef TPP_MLIR
registry.insert<mlir::xsmm::XsmmDialect>();
registry.insert<mlir::check::CheckDialect>();
registry.insert<mlir::perf::PerfDialect>();
#endif
}

registerAllDialects(registry);
registerAllExtensions(registry);
Expand All @@ -339,5 +358,20 @@ MLIRContext* get_shared_mlir_context() {
} // namespace

void ov::pass::transformMLIR(std::shared_ptr<ov::Model> model) {
injectMLIR(model, get_shared_mlir_context());
if(util::getenv_bool("OV_MLIR", true)) {
bool tpp_mlir_default =
#ifdef TPP_MLIR
true;
#else
false;
#endif
bool tpp_mlir_enabled = util::getenv_bool("OV_MLIR_TPP", tpp_mlir_default);
#ifndef TPP_MLIR
OPENVINO_ASSERT(!tpp_mlir_enabled,
"[ ERROR ] OpenVINO wasn't compiled with TPP_MLIR support, "
"but OV_MLIR_TPP environment variable is set to enable it.");
#endif

injectMLIR(model, get_shared_mlir_context(tpp_mlir_enabled), tpp_mlir_enabled);
}
}
6 changes: 6 additions & 0 deletions src/common/transformations/src/transformations/mlir/convert_common.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,9 @@

#include "convert_common.hpp"

#include <openvino/util/env_util.hpp>


namespace {

using namespace mlir;
Expand Down Expand Up @@ -58,6 +61,9 @@ IntegerType getBool8Type(MLIRContext* ctx) {
namespace ov {
namespace mlir {

bool is_debug() {
util::getenv_bool("OV_MLIR_DEBUG", false);
}

Location createLayerLocation(MLIRContext* ctx, const std::string& layerName, const std::string& layerType) {
const auto layerNameAttr = StringAttr::get(ctx, layerName);
Expand Down
7 changes: 7 additions & 0 deletions src/common/transformations/src/transformations/mlir/convert_common.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@

#pragma once

#include <iostream>

#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/MLIRContext.h"
Expand All @@ -17,6 +19,11 @@
namespace ov {
namespace mlir {

bool is_debug();

#define OPENVINO_MLIR_DEBUG(X) do if(::ov::mlir::is_debug()) { X; } while(false)
#define OPENVINO_MLIR_DEBUG_PRINT(X) do if(::ov::mlir::is_debug()) { ::std::cerr << X; } while(false)

using namespace ::mlir;

Location createLayerLocation(MLIRContext* ctx, const std::string& layerName, const std::string& layerType);
Expand Down
Loading

0 comments on commit c793696

Please sign in to comment.