diff --git a/llvm/src/lifting/function_to_sdfg.cpp b/llvm/src/lifting/function_to_sdfg.cpp index 451488694..2e402b06d 100644 --- a/llvm/src/lifting/function_to_sdfg.cpp +++ b/llvm/src/lifting/function_to_sdfg.cpp @@ -702,7 +702,7 @@ std::unique_ptr FunctionToSDFG::simplify(std::unique_ptr #include +#include "sdfg/passes/schedules/expansion_pass.h" + int main(int argc, char* argv[]) { sdfg::codegen::register_default_dispatchers(); sdfg::serializer::register_default_serializers(); @@ -29,8 +31,8 @@ int main(int argc, char* argv[]) { sdfg::builder::StructuredSDFGBuilder builder(*sdfg); sdfg::analysis::AnalysisManager analysis_manager(builder.subject()); - sdfg::passes::Pipeline libnode_expansion = sdfg::passes::Pipeline::expansion(); - libnode_expansion.run(builder, analysis_manager); + sdfg::passes::MathExpansionPass math_expansion; + math_expansion.run(builder, analysis_manager); sdfg::passes::TensorToPointerConversionPass tensor_to_pointer_conversion_pass; tensor_to_pointer_conversion_pass.run(builder, analysis_manager); diff --git a/python/bindings/bindings.cpp b/python/bindings/bindings.cpp index 17cb1a09d..65ff24cb5 100644 --- a/python/bindings/bindings.cpp +++ b/python/bindings/bindings.cpp @@ -229,6 +229,12 @@ PYBIND11_MODULE(_sdfg, m) { py::arg("dump_json") = true, py::arg("record_for_instrumentation") = false ) + .def( + "set_output_dir", + static_cast(&PyStructuredSDFG::set_output_dir), + py::arg("path"), + "Set the output directory" + ) .def("normalize", &PyStructuredSDFG::normalize, "Normalize the SDFG") .def( "schedule", diff --git a/python/bindings/py_structured_sdfg.cpp b/python/bindings/py_structured_sdfg.cpp index d5f3457b8..a98be5dec 100644 --- a/python/bindings/py_structured_sdfg.cpp +++ b/python/bindings/py_structured_sdfg.cpp @@ -62,6 +62,7 @@ #include "sdfg/passes/rpc/daisytuner_rpc_context.h" #include "sdfg/passes/rpc/rpc_context.h" #include "sdfg/passes/rpc/rpc_scheduler.h" +#include "sdfg/passes/schedules/expansion_pass.h" #include "sdfg/passes/targets/target_mapping_pass.h" #include "sdfg/util/offloading_instrumentation_plan.h" #include "targets/target_mapping.h" @@ -113,6 +114,10 @@ PyStructuredSDFG PyStructuredSDFG::from_sdfg(sdfg::plugins::Context& ctx, std::u std::string PyStructuredSDFG::name() const { return sdfg_->name(); } +void PyStructuredSDFG::set_output_dir(const std::string& path) { this->set_output_dir(std::filesystem::path(path)); } + +void PyStructuredSDFG::set_output_dir(const std::filesystem::path& path) { this->output_dir_ = path; } + sdfg::plugins::Context& PyStructuredSDFG::docc_context() const { return docc_context_; } const sdfg::types::IType& PyStructuredSDFG::return_type() const { return sdfg_->return_type(); } @@ -182,8 +187,8 @@ void PyStructuredSDFG::expand(const docc::target::TargetOptions& options) { local_buffer_reuse_pipeline.run(builder_opt, analysis_manager); // Expand library nodes - sdfg::passes::Pipeline libnode_expansion = sdfg::passes::Pipeline::expansion(); - libnode_expansion.run(builder_opt, analysis_manager); + sdfg::passes::MathExpansionPass math_expand; + math_expand.run(builder_opt, analysis_manager); sdfg::passes::TensorToPointerConversionPass tensor_to_pointer_conversion_pass; tensor_to_pointer_conversion_pass.run(builder_opt, analysis_manager); @@ -448,10 +453,7 @@ std::string PyStructuredSDFG::compile( sdfg::analysis::AnalysisManager analysis_manager(*sdfg_); - // Run expansion pass - sdfg::passes::Pipeline expansion = sdfg::passes::Pipeline::expansion(); sdfg::builder::StructuredSDFGBuilder builder_opt(*sdfg_); - expansion.run(builder_opt, analysis_manager); // Instrumentation plan std::unique_ptr instrumentation_plan; diff --git a/python/bindings/py_structured_sdfg.h b/python/bindings/py_structured_sdfg.h index 3e728e12c..20fefe129 100644 --- a/python/bindings/py_structured_sdfg.h +++ b/python/bindings/py_structured_sdfg.h @@ -16,6 +16,7 @@ class PyStructuredSDFG { private: sdfg::plugins::Context& docc_context_; std::unique_ptr sdfg_; + std::optional output_dir_; PyStructuredSDFG(sdfg::plugins::Context& ctx, std::unique_ptr& sdfg); @@ -34,6 +35,9 @@ class PyStructuredSDFG { std::string name() const; + void set_output_dir(const std::string& path); + void set_output_dir(const std::filesystem::path& path); + sdfg::plugins::Context& docc_context() const; sdfg::StructuredSDFG& sdfg() { return *sdfg_; } diff --git a/python/docc/compiler/docc_program.py b/python/docc/compiler/docc_program.py index efb15701e..26d9ad5d1 100644 --- a/python/docc/compiler/docc_program.py +++ b/python/docc/compiler/docc_program.py @@ -157,6 +157,8 @@ def sdfg_pipe( target_options.target = self.target target_options.category = self.category target_options.remote_tuning = remote_tuning + if self.debug_dump: + sdfg.set_output_dir(output_folder) # Einsum detection sdfg.einsum() diff --git a/sdfg/include/sdfg/data_flow/library_nodes/math/tensor/reduce_node.h b/sdfg/include/sdfg/data_flow/library_nodes/math/tensor/reduce_node.h index 544aeaa10..56c810085 100644 --- a/sdfg/include/sdfg/data_flow/library_nodes/math/tensor/reduce_node.h +++ b/sdfg/include/sdfg/data_flow/library_nodes/math/tensor/reduce_node.h @@ -191,6 +191,8 @@ class ReduceNode : public TensorNode { data_flow::PointerAccessType pointer_access_type(int input_idx) const override; + std::string toStr() const override; + protected: virtual bool expand_inner( builder::StructuredSDFGBuilder& builder, diff --git a/sdfg/include/sdfg/data_flow/library_nodes/math/tensor/tensor_layout.h b/sdfg/include/sdfg/data_flow/library_nodes/math/tensor/tensor_layout.h index ee0c95a7c..777acbdc9 100644 --- a/sdfg/include/sdfg/data_flow/library_nodes/math/tensor/tensor_layout.h +++ b/sdfg/include/sdfg/data_flow/library_nodes/math/tensor/tensor_layout.h @@ -101,6 +101,8 @@ class TensorLayout { std::unique_ptr squeeze() const; std::unique_ptr reshape(const symbolic::MultiExpression& new_shape) const; + + static std::ostream& emit_symbolic_list(std::ostream& stream, const symbolic::MultiExpression& list); }; std::ostream& operator<<(std::ostream& stream, const TensorLayout& layout); diff --git a/sdfg/include/sdfg/passes/pipeline.h b/sdfg/include/sdfg/passes/pipeline.h index a3b2fed75..82b491b9c 100644 --- a/sdfg/include/sdfg/passes/pipeline.h +++ b/sdfg/include/sdfg/passes/pipeline.h @@ -56,8 +56,6 @@ class Pipeline : public Pass { static Pipeline data_parallelism(); static Pipeline memory(); - - static Pipeline expansion(); }; } // namespace passes diff --git a/sdfg/include/sdfg/passes/schedules/expansion_pass.h b/sdfg/include/sdfg/passes/schedules/expansion_pass.h index e527ac296..729c0393b 100644 --- a/sdfg/include/sdfg/passes/schedules/expansion_pass.h +++ b/sdfg/include/sdfg/passes/schedules/expansion_pass.h @@ -23,52 +23,49 @@ #pragma once +#include "sdfg/data_flow/library_nodes/math/math_node.h" #include "sdfg/passes/pass.h" #include "sdfg/visitor/structured_sdfg_visitor.h" namespace sdfg { namespace passes { -/** - * @class Expansion - * @brief Visitor that expands library nodes into primitive operations - * - * The Expansion visitor traverses the SDFG and expands library nodes that - * have ImplementationType_NONE. This allows high-level mathematical operations - * to be transformed into lower-level constructs that can be optimized and - * scheduled. - */ -class Expansion : public visitor::StructuredSDFGVisitor { +class MathExpansionPass; + +class MathExpansionVisitor : public visitor::ActualStructuredSDFGVisitor { + friend MathExpansionPass; + +private: + builder::StructuredSDFGBuilder& builder_; + analysis::AnalysisManager& analysis_manager_; + + struct LibNodeContainer { + math::MathNode& node; + structured_control_flow::Block& block; + }; + + std::vector nodes_to_expand_; + public: /** * @brief Construct the expansion visitor * @param builder SDFG builder for creating new nodes * @param analysis_manager Analysis manager for querying properties */ - Expansion(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager); + MathExpansionVisitor(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager); - /** - * @brief Get the pass name - * @return Name of the pass - */ - static std::string name() { return "Expansion"; }; - - /** - * @brief Visit a block and attempt to expand its library nodes - * @param node Block to visit - * @return True if any expansion occurred - */ - bool accept(structured_control_flow::Block& node) override; + bool visit(sdfg::structured_control_flow::Block& node) override; }; /** - * @typedef ExpansionPass - * @brief Pass wrapper for the Expansion visitor - * - * This typedef creates a pass from the Expansion visitor, allowing it to be - * used in the pass pipeline system. + * @class MathExpansionPass + * @brief Looks for and expands math-nodes that are not already mapped to a specific target */ -typedef VisitorPass ExpansionPass; +class MathExpansionPass : public Pass { + std::string name() override { return "MathExpansion"; } + + bool run_pass(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) override; +}; } // namespace passes } // namespace sdfg diff --git a/sdfg/src/data_flow/library_nodes/math/tensor/reduce_node.cpp b/sdfg/src/data_flow/library_nodes/math/tensor/reduce_node.cpp index e19c53de8..4d99e2e1d 100644 --- a/sdfg/src/data_flow/library_nodes/math/tensor/reduce_node.cpp +++ b/sdfg/src/data_flow/library_nodes/math/tensor/reduce_node.cpp @@ -100,6 +100,27 @@ data_flow::PointerAccessType ReduceNode::pointer_access_type(int input_idx) cons } } +std::ostream& operator<<(std::ostream& os, const std::vector& list) { + os << "["; + for (size_t i = 0; i < list.size(); ++i) { + if (i > 0) os << ", "; + os << list[i]; + } + os << "]"; + return os; +} + +std::string ReduceNode::toStr() const { + std::stringstream ss; + ss << this->code_.value(); + ss << "(shape="; + TensorLayout::emit_symbolic_list(ss, shape_); + ss << ", axes=" << axes_; + ss << ", keep=" << this->keepdims_; + ss << ")"; + return ss.str(); +} + bool ReduceNode::expand_inner( builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager, diff --git a/sdfg/src/data_flow/library_nodes/math/tensor/tensor_layout.cpp b/sdfg/src/data_flow/library_nodes/math/tensor/tensor_layout.cpp index 96dffb9de..373661709 100644 --- a/sdfg/src/data_flow/library_nodes/math/tensor/tensor_layout.cpp +++ b/sdfg/src/data_flow/library_nodes/math/tensor/tensor_layout.cpp @@ -107,18 +107,21 @@ TensorLayout TensorLayout::deserialize_from_json(const nlohmann::json& j) { return std::move(TensorLayout(shape, strides, offset)); } -std::ostream& operator<<(std::ostream& stream, const TensorLayout& layout) { - stream << "{shape["; - for (size_t i = 0; i < layout.shape().size(); ++i) { - if (i > 0) stream << ", "; - stream << layout.shape().at(i)->__str__(); - } - stream << "], strides=["; - for (size_t i = 0; i < layout.strides().size(); ++i) { +std::ostream& TensorLayout::emit_symbolic_list(std::ostream& stream, const symbolic::MultiExpression& list) { + stream << "["; + for (size_t i = 0; i < list.size(); ++i) { if (i > 0) stream << ", "; - stream << layout.strides().at(i)->__str__(); + stream << list.at(i)->__str__(); } stream << "]"; + return stream; +} + +std::ostream& operator<<(std::ostream& stream, const TensorLayout& layout) { + stream << "{shape="; + TensorLayout::emit_symbolic_list(stream, layout.shape()); + stream << ", strides="; + TensorLayout::emit_symbolic_list(stream, layout.strides()); if (SymEngine::neq(*layout.offset(), *symbolic::integer(0))) { stream << ", off=" << layout.offset()->__str__(); } diff --git a/sdfg/src/passes/pipeline.cpp b/sdfg/src/passes/pipeline.cpp index ddd8dc428..24e4d6f66 100644 --- a/sdfg/src/passes/pipeline.cpp +++ b/sdfg/src/passes/pipeline.cpp @@ -173,13 +173,5 @@ Pipeline Pipeline::memory() { return p; }; -Pipeline Pipeline::expansion() { - Pipeline p("Expansion"); - - p.register_pass(); - - return p; -}; - } // namespace passes } // namespace sdfg diff --git a/sdfg/src/passes/schedules/expansion_pass.cpp b/sdfg/src/passes/schedules/expansion_pass.cpp index 6edcd2bac..72b6239d1 100644 --- a/sdfg/src/passes/schedules/expansion_pass.cpp +++ b/sdfg/src/passes/schedules/expansion_pass.cpp @@ -5,26 +5,48 @@ namespace sdfg { namespace passes { -Expansion::Expansion(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) - : visitor::StructuredSDFGVisitor(builder, analysis_manager) {} +MathExpansionVisitor:: + MathExpansionVisitor(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) + : visitor::ActualStructuredSDFGVisitor(), builder_(builder), analysis_manager_(analysis_manager) {} -bool Expansion::accept(structured_control_flow::Block& node) { +bool MathExpansionVisitor::visit(structured_control_flow::Block& node) { auto& dataflow = node.dataflow(); - bool applied = false; + for (auto* library_node : dataflow.library_nodes()) { if (library_node->implementation_type() != data_flow::ImplementationType_NONE) { continue; } if (auto math_node = dynamic_cast(library_node)) { - if (math_node->expand(this->builder_, this->analysis_manager_)) { - return true; - } + this->nodes_to_expand_.emplace_back(*math_node, node); } } - return applied; + return true; } +bool MathExpansionPass::run_pass(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) { + MathExpansionVisitor v(builder, analysis_manager); + + v.dispatch(builder.subject().root()); + + auto& nodes = v.nodes_to_expand_; + + bool expanded_any = false; + + for (auto& entry : std::views::reverse(nodes)) { + // TODO: check if the prerequisits are met, like if the libNode is standalone or if we need to cut it out of a + // larger block first + + if (entry.node.expand(builder, analysis_manager)) { + // If expansion was successful, remove the original library node // TODO requires new API to do this clean + // builder.remove_node(entry.block, entry.node); + // remove block + expanded_any |= true; + } + } + + return expanded_any; +} } // namespace passes } // namespace sdfg diff --git a/sdfg/tests/passes/schedules/expansion_pass_test.cpp b/sdfg/tests/passes/schedules/expansion_pass_test.cpp index 2552dba67..0f8e0be03 100644 --- a/sdfg/tests/passes/schedules/expansion_pass_test.cpp +++ b/sdfg/tests/passes/schedules/expansion_pass_test.cpp @@ -47,7 +47,7 @@ TEST(ExpansionPassTest, MeanNode_2D) { EXPECT_EQ(block.dataflow().nodes().size(), 3); analysis::AnalysisManager analysis_manager(builder.subject()); - passes::ExpansionPass expansion_pass; + passes::MathExpansionPass expansion_pass; EXPECT_TRUE(expansion_pass.run(builder, analysis_manager)); dump_sdfg(builder.subject(), "1.expanded"); @@ -86,7 +86,7 @@ TEST(ExpansionPassTest, StdNode_1D) { dump_sdfg(builder.subject(), "0.init"); analysis::AnalysisManager analysis_manager(builder.subject()); - passes::ExpansionPass expansion_pass; + passes::MathExpansionPass expansion_pass; EXPECT_TRUE(expansion_pass.run(builder, analysis_manager)); dump_sdfg(builder.subject(), "1.expanded");