Skip to content

Commit

Permalink
fix test
Browse files Browse the repository at this point in the history
  • Loading branch information
PhilippGrulich committed Oct 15, 2024
1 parent f65b1d7 commit 6ad1bb3
Show file tree
Hide file tree
Showing 3 changed files with 74 additions and 46 deletions.
7 changes: 5 additions & 2 deletions nautilus/src/nautilus/compiler/JITCompiler.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
#include "nautilus/compiler/ir/util/GraphVizUtil.hpp"
#include "nautilus/config.hpp"
#include "nautilus/exceptions/RuntimeException.hpp"
#include "nautilus/logging.hpp"
#include <chrono>
#include <iomanip>
#include <iostream>
Expand Down Expand Up @@ -78,8 +79,10 @@ std::unique_ptr<Executable> JITCompiler::compile(JITCompiler::wrapper_function f
auto irGenerationPhase = tracing::TraceToIRConversionPhase();
auto ir = irGenerationPhase.apply(std::move(afterSSA), compilationId);
dumpHandler.dump("after_ir_creation", "ir", [&]() { return ir->toString(); });
auto url = ir::createGraphVizFromIr(ir);
std::cout << url << std::endl;
if (options.getOptionOrDefault("dump.graph", false)) {
auto url = ir::createGraphVizFromIr(ir);
log::error("{}", url);
}
// lower to backend
auto backendName = options.getOptionOrDefault<std::string>("engine.backend", "mlir");
auto backend = backends->getBackend(backendName);
Expand Down
112 changes: 68 additions & 44 deletions nautilus/src/nautilus/compiler/ir/util/GraphVizUtil.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -241,56 +241,75 @@ class GraphvizWriter {
getBlockArgumentMap(branchOp->getNextBlockInvocation().getBlock(), visitedBlocks, map);
}
}
void write_edges(const std::shared_ptr<IRGraph>& graph, [[maybe_unused]] bool hideIntermediateBlockArguments) {
void write_edges(const std::shared_ptr<IRGraph>& graph, bool hideIntermediateBlockArguments, bool writeBlocksOnly) {

std::map<Operation*, std::vector<Operation*>> operatorDataflowMapping;
auto& rootBlock = graph->getRootOperation().getBasicBlocks()[0];
std::set<const BasicBlock*> vistedBlocks;
getBlockArgumentMap(rootBlock.get(), vistedBlocks, operatorDataflowMapping);
// crate dataflow edges
for (const auto& block : graph->getRootOperation().getBasicBlocks()) {
for (const auto& op : block->getOperations()) {

const auto& to = nodeIdMap[op.get()];
// process inputs of node
for (auto input : op->getInputs()) {
const auto& from = nodeIdMap[input];
write_edge(from, to, "data", input->getIdentifier().toString());
}

// handle control-flow arguments
if (writeBlocksOnly) {
auto& op = block->getOperations().back();
if (op->getOperationType() == Operation::OperationType::IfOp) {
[[maybe_unused]] auto ifOp = as<IfOperation>(op);
write_block_argument_edges(graph, ifOp->getTrueBlockInvocation(), hideIntermediateBlockArguments);
write_block_argument_edges(graph, ifOp->getFalseBlockInvocation(), hideIntermediateBlockArguments);
auto ifOp = as<IfOperation>(op);
auto trueBlock = ifOp->getTrueBlockInvocation().getBlock()->getIdentifier();
write_edge(block->getIdentifier(), trueBlock, "control");
auto falseBlock = ifOp->getFalseBlockInvocation().getBlock()->getIdentifier();
write_edge(block->getIdentifier(), falseBlock, "control");
} else if (op->getOperationType() == Operation::OperationType::BranchOp) {
[[maybe_unused]] auto branchOp = as<BranchOperation>(op);
write_block_argument_edges(graph, branchOp->getNextBlockInvocation(),
hideIntermediateBlockArguments);
auto branchOp = as<BranchOperation>(op);
auto falseBlock = branchOp->getNextBlockInvocation().getBlock()->getIdentifier();
write_edge(block->getIdentifier(), falseBlock, "control");
}
}
}

// create control-flow edges
for (const auto& blocks : graph->getRootOperation().getBasicBlocks()) {
std::string from = "start_" + blocks->getIdentifier();
for (const auto& op : blocks->getOperations()) {
if (isControlFlowOp(op.get())) {
} else {
for (const auto& op : block->getOperations()) {
const auto& to = nodeIdMap[op.get()];
write_edge(from, to, "control");
from = to;
// process inputs of node
for (auto input : op->getInputs()) {
const auto& from = nodeIdMap[input];
write_edge(from, to, "data", input->getIdentifier().toString());
}

// handle control-flow arguments
if (op->getOperationType() == Operation::OperationType::IfOp) {
[[maybe_unused]] auto ifOp = as<IfOperation>(op);
write_block_argument_edges(graph, ifOp->getTrueBlockInvocation(),
hideIntermediateBlockArguments);
write_block_argument_edges(graph, ifOp->getFalseBlockInvocation(),
hideIntermediateBlockArguments);
} else if (op->getOperationType() == Operation::OperationType::BranchOp) {
[[maybe_unused]] auto branchOp = as<BranchOperation>(op);
write_block_argument_edges(graph, branchOp->getNextBlockInvocation(),
hideIntermediateBlockArguments);
}
}
}
}

if (op->getOperationType() == Operation::OperationType::IfOp) {
auto ifOp = as<IfOperation>(op);
std::string toTrue = "start_" + ifOp->getTrueBlockInvocation().getBlock()->getIdentifier();
write_edge(from, toTrue, "control", "true");
std::string toFalse = "start_" + ifOp->getFalseBlockInvocation().getBlock()->getIdentifier();
write_edge(from, toFalse, "control", "false");
} else if (op->getOperationType() == Operation::OperationType::BranchOp) {
auto branchOp = as<BranchOperation>(op);
std::string to = "start_" + branchOp->getNextBlockInvocation().getBlock()->getIdentifier();
write_edge(from, to, "control");
if (!writeBlocksOnly) {

// create control-flow edges
for (const auto& blocks : graph->getRootOperation().getBasicBlocks()) {
std::string from = "start_" + blocks->getIdentifier();
for (const auto& op : blocks->getOperations()) {
if (isControlFlowOp(op.get())) {
const auto& to = nodeIdMap[op.get()];
write_edge(from, to, "control");
from = to;
}

if (op->getOperationType() == Operation::OperationType::IfOp) {
auto ifOp = as<IfOperation>(op);
std::string toTrue = "start_" + ifOp->getTrueBlockInvocation().getBlock()->getIdentifier();
write_edge(from, toTrue, "control", "true");
std::string toFalse = "start_" + ifOp->getFalseBlockInvocation().getBlock()->getIdentifier();
write_edge(from, toFalse, "control", "false");
} else if (op->getOperationType() == Operation::OperationType::BranchOp) {
auto branchOp = as<BranchOperation>(op);
std::string to = "start_" + branchOp->getNextBlockInvocation().getBlock()->getIdentifier();
write_edge(from, to, "control");
}
}
}
}
Expand Down Expand Up @@ -323,15 +342,16 @@ class GraphvizWriter {
}
stream << " node" << from << " -> node" << to << " " << write_attrs(attrs) << ";" << std::endl;
}
void write_graph(const std::shared_ptr<IRGraph>& graph, bool hidpi = false, bool draw_blocks = true) {
void write_graph(const std::shared_ptr<IRGraph>& graph, bool hidpi = false, bool draw_blocks = true,
bool drawBlocksOnly = true) {
auto attrs = std::map<std::string, std::string>();
attrs["bgcolor"] = "white";
if (hidpi) {
attrs["dpi"] = "200";
}
start_graph(attrs);
write_nodes(graph, draw_blocks);
write_edges(graph, true);
write_nodes(graph, draw_blocks, drawBlocksOnly);
write_edges(graph, true, drawBlocksOnly);
end_graph();
}
void write_node(const std::string& indent, const std::string& label, const std::string& id,
Expand Down Expand Up @@ -381,14 +401,18 @@ class GraphvizWriter {
write_node_for_op(indent, block->getIdentifier(), op.get());
}
}
void write_nodes(const std::shared_ptr<IRGraph>& graph, [[maybe_unused]] bool draw_blocks) {
void write_nodes(const std::shared_ptr<IRGraph>& graph, bool draw_blocks, bool drawBlocksOnly) {

if (draw_blocks) {
// Assuming you have a way to divide nodes into blocks
for (const auto& block : graph->getRootOperation().getBasicBlocks()) {
start_subgraph(block->getIdentifier());
write_nodes_for_block(block);
end_subgraph();
if (!drawBlocksOnly) {
start_subgraph(block->getIdentifier());
write_nodes_for_block(block);
end_subgraph();
} else {
write_node(" ", "Block" + block->getIdentifier(), block->getIdentifier(), "control");
}
}
} else {
for (const auto& block : graph->getRootOperation().getBasicBlocks()) {
Expand Down
1 change: 1 addition & 0 deletions nautilus/test/execution-tests/ExecutionTest.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1000,6 +1000,7 @@ TEST_CASE("Engine Compiler Test") {
engine::Options options;
options.setOption("engine.backend", backend);
options.setOption("dump.all", backend);
options.setOption("dump.graph", true);
auto engine = engine::NautilusEngine(options);
runAllTests(engine);
}
Expand Down

0 comments on commit 6ad1bb3

Please sign in to comment.