29989607664

Committed 23 Jul 2026 07:54AM UTC coverage: 64.081% (+0.3%) from 63.787%

Build # 29989607664

Build Type

push

github

Committed by

web-flow

Commit Message

Add tensor concatenation to PyTorch frontend (and MLIR frontend) (#868)

- Added a concat tensor node that concatenates arbitrarily many tensors
- Mapped `torch.ops.aten.cat.default` to concat tensor node
- Changed translation of `tensor.concat` MLIR op to concat tensor node instead of loop nest
- Added replace_symbols method to all types to replace symbols in types like the tensor type
- Added target specific expansion for concat tensor node to avoid branching on GPUs

Coverage Stats

352 of 454 new or added lines in 19 files covered. (77.53%)

1 existing line in 1 file now uncovered.

42653 of 66561 relevant lines covered (64.08%)

739.89 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

53.33

/opt/src/passes/offloading/cuda_library_node_expansion_pass.cpp

#include "sdfg/passes/offloading/cuda_library_node_expansion_pass.h"

#include "sdfg/data_flow/library_nodes/math/tensor/concat_node.h"
#include "sdfg/data_flow/library_nodes/math/tensor/conv_node.h"
#include "sdfg/data_flow/library_nodes/math/tensor/matmul_node.h"
#include "sdfg/data_flow/library_nodes/math/tensor/reduce_ops/softmax_node.h"
#include "sdfg/targets/cuda/cuda.h"
#include "sdfg/targets/cuda/math/tensor/batched_matmul_expander.h"
#include "sdfg/targets/cuda/math/tensor/concat_expander.h"
#include "sdfg/targets/cuda/math/tensor/conv_expander.h"

namespace sdfg {
namespace passes {

CudaExpansion::CudaExpansion(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager)
    : visitor::NonStoppingStructuredSDFGVisitor(builder, analysis_manager) {}

bool CudaExpansion::accept(structured_control_flow::Block& node) {
    auto& dataflow = node.dataflow();

    bool made_changes = false;

    for (auto* library_node : dataflow.library_nodes()) {
        if (library_node->implementation_type() != data_flow::ImplementationType_NONE) {
            continue;
        }

        auto& lib_node_code = library_node->code();

        if (lib_node_code == math::tensor::LibraryNodeType_Conv) {
            auto& conv_node = static_cast<math::tensor::ConvNode&>(*library_node);
            sdfg::offloading::CudaConvExpander expander(conv_node);
            made_changes |= expander.expand(builder_, analysis_manager_);
        } else if (lib_node_code == math::tensor::LibraryNodeType_MatMul) {
            auto& matmul_node = static_cast<math::tensor::MatMulNode&>(*library_node);
            sdfg::offloading::CudaBatchedMatMulExpander expander(matmul_node);
            made_changes |= expander.expand(builder_, analysis_manager_);
        } else if (lib_node_code == math::tensor::LibraryNodeType_Softmax) {
            library_node->implementation_type() = cuda::ImplementationType_CUDAWithTransfers;
            made_changes = true;
        } else if (lib_node_code == math::tensor::LibraryNodeType_TensorConcat) {
            auto& concat_node = static_cast<math::tensor::ConcatNode&>(*library_node);
            sdfg::offloading::CudaConcatExpander expander(concat_node);
            made_changes |= expander.expand(builder_, analysis_manager_);
        } else {
            continue;
        }
    }
    return made_changes;
};

} // namespace passes
} // namespace sdfg

1	#include "sdfg/passes/offloading/cuda_library_node_expansion_pass.h"
2
3	#include "sdfg/data_flow/library_nodes/math/tensor/concat_node.h"
4	#include "sdfg/data_flow/library_nodes/math/tensor/conv_node.h"
5	#include "sdfg/data_flow/library_nodes/math/tensor/matmul_node.h"
6	#include "sdfg/data_flow/library_nodes/math/tensor/reduce_ops/softmax_node.h"
7	#include "sdfg/targets/cuda/cuda.h"
8	#include "sdfg/targets/cuda/math/tensor/batched_matmul_expander.h"
9	#include "sdfg/targets/cuda/math/tensor/concat_expander.h"
10	#include "sdfg/targets/cuda/math/tensor/conv_expander.h"
11
12	namespace sdfg {
13	namespace passes {
14
15	CudaExpansion::CudaExpansion(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager)
16	: visitor::NonStoppingStructuredSDFGVisitor(builder, analysis_manager) {}	1✔
17
18	bool CudaExpansion::accept(structured_control_flow::Block& node) {	1✔
19	auto& dataflow = node.dataflow();	1✔
20
21	bool made_changes = false;	1✔
22
23	for (auto* library_node : dataflow.library_nodes()) {	1✔
24	if (library_node->implementation_type() != data_flow::ImplementationType_NONE) {	1✔
25	continue;	×
26	}	×
27
28	auto& lib_node_code = library_node->code();	1✔
29
30	if (lib_node_code == math::tensor::LibraryNodeType_Conv) {	1✔
31	auto& conv_node = static_cast<math::tensor::ConvNode&>(*library_node);	×
32	sdfg::offloading::CudaConvExpander expander(conv_node);	×
33	made_changes \|= expander.expand(builder_, analysis_manager_);	×
34	} else if (lib_node_code == math::tensor::LibraryNodeType_MatMul) {	1✔
35	auto& matmul_node = static_cast<math::tensor::MatMulNode&>(*library_node);	×
36	sdfg::offloading::CudaBatchedMatMulExpander expander(matmul_node);	×
37	made_changes \|= expander.expand(builder_, analysis_manager_);	×
38	} else if (lib_node_code == math::tensor::LibraryNodeType_Softmax) {	1✔
39	library_node->implementation_type() = cuda::ImplementationType_CUDAWithTransfers;	1✔
40	made_changes = true;	1✔
41	} else if (lib_node_code == math::tensor::LibraryNodeType_TensorConcat) {	1✔
NEW 42	auto& concat_node = static_cast<math::tensor::ConcatNode&>(*library_node);	×
NEW 43	sdfg::offloading::CudaConcatExpander expander(concat_node);	×
NEW 44	made_changes \|= expander.expand(builder_, analysis_manager_);	×
45	} else {	×
46	continue;	×
47	}	×
48	}	1✔
49	return made_changes;	1✔
50	};	1✔
51
52	} // namespace passes
53	} // namespace sdfg

daisytuner / docc / 29989607664

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous