From 3b30f3b7e7c7e7f3d4aa55b9db442810540d193c Mon Sep 17 00:00:00 2001 From: Ramon Wirsch Date: Mon, 8 Jun 2026 13:23:07 +0200 Subject: [PATCH 1/3] Switched expansion "pipeline" to MathExpansionPass. * Will collect all math-nodes and attempt to expand them all in reverse execution order to minimize effects of expands on other analysis ! does not clear any analysis (which makes it fragile. Should not matter for current expands(), but slightly dangerous until we have removed the need for ScopeAnalysis or can more granularly clear analyses. --- llvm/src/lifting/function_to_sdfg.cpp | 2 +- python/bindings/py_structured_sdfg.cpp | 8 +-- sdfg/include/sdfg/passes/pipeline.h | 2 - .../sdfg/passes/schedules/expansion_pass.h | 55 +++++++++---------- sdfg/src/passes/pipeline.cpp | 8 --- sdfg/src/passes/schedules/expansion_pass.cpp | 38 ++++++++++--- .../passes/schedules/expansion_pass_test.cpp | 4 +- 7 files changed, 62 insertions(+), 55 deletions(-) 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 instrumentation_plan; 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/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"); From 873c6a95811d4c1559e70ea1f84187a666c60204 Mon Sep 17 00:00:00 2001 From: Ramon Wirsch Date: Mon, 8 Jun 2026 13:29:24 +0200 Subject: [PATCH 2/3] Also migrated sdfg-json-to-c.cpp to not using the expansion pipeline anymore --- mlir/test/sdfg-json-to-c/sdfg-json-to-c.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/mlir/test/sdfg-json-to-c/sdfg-json-to-c.cpp b/mlir/test/sdfg-json-to-c/sdfg-json-to-c.cpp index fc3d33afd..7bf8b4041 100644 --- a/mlir/test/sdfg-json-to-c/sdfg-json-to-c.cpp +++ b/mlir/test/sdfg-json-to-c/sdfg-json-to-c.cpp @@ -12,6 +12,8 @@ #include #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); From 37383f1fcd6a8ec4c32f067bba4431408c63294f Mon Sep 17 00:00:00 2001 From: Ramon Wirsch Date: Mon, 8 Jun 2026 16:08:31 +0200 Subject: [PATCH 3/3] toStr() for ReduceNodes and giving PyStructuredSDFG access to the output_dir for additional dumping --- python/bindings/bindings.cpp | 6 ++++++ python/bindings/py_structured_sdfg.cpp | 4 ++++ python/bindings/py_structured_sdfg.h | 4 ++++ python/docc/compiler/docc_program.py | 2 ++ .../library_nodes/math/tensor/reduce_node.h | 2 ++ .../library_nodes/math/tensor/tensor_layout.h | 2 ++ .../library_nodes/math/tensor/reduce_node.cpp | 21 +++++++++++++++++++ .../math/tensor/tensor_layout.cpp | 21 +++++++++++-------- 8 files changed, 53 insertions(+), 9 deletions(-) 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 69bdfc570..a98be5dec 100644 --- a/python/bindings/py_structured_sdfg.cpp +++ b/python/bindings/py_structured_sdfg.cpp @@ -114,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(); } 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/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__(); }