30338860029

Committed 28 Jul 2026 07:33AM UTC coverage: 64.09% (-0.05%) from 64.141%

Build # 30338860029

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #893 from daisytuner/fix-softmax-dispatcher

Fix softmax dispatcher for strided accesses

Coverage Stats

36 of 100 new or added lines in 1 file covered. (36.0%)

42867 of 66886 relevant lines covered (64.09%)

731.09 hits per line

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

74.53

/opt/src/targets/cuda/math/tensor/softmax.cpp

#include "sdfg/targets/cuda/math/tensor/softmax.h"
#include "sdfg/symbolic/symbolic.h"
#include "sdfg/targets/cuda/cuda.h"

#include <algorithm>

namespace sdfg::cuda::tensor {

static constexpr int SOFTMAX_BLOCK_SIZE = 256;

// Contiguous softmax kernel: one block per row, warp-shuffle + shared-memory reductions.
// Used when the reduced axis is innermost (inner == 1), so a row's elements are contiguous
// in memory and threads within a warp access consecutive addresses (coalesced).
static void
emit_softmax_kernel_contiguous(codegen::PrettyPrinter& ks, const std::string& kernel_name, const std::string& type) {
    ks << "__global__ void " << kernel_name << "(const " << type << "* __restrict__ input, " << type
       << "* __restrict__ output, int num_rows, int row_size) {" << std::endl;
    ks.setIndent(ks.indent() + 4);

    ks << "int row = blockIdx.x;" << std::endl;
    ks << "if (row >= num_rows) return;" << std::endl;
    ks << std::endl;
    ks << "const " << type << "* row_in = input + (size_t)row * row_size;" << std::endl;
    ks << type << "* row_out = output + (size_t)row * row_size;" << std::endl;
    ks << std::endl;

    // Shared memory for cross-warp reduction
    ks << "extern __shared__ " << type << " sdata[];" << std::endl;
    ks << "int lane_id = threadIdx.x & 31;" << std::endl;
    ks << "int warp_id = threadIdx.x >> 5;" << std::endl;
    ks << "int num_warps = (blockDim.x + 31) >> 5;" << std::endl;
    ks << std::endl;

    // Phase 1: find row max
    ks << "// Phase 1: row max" << std::endl;
    ks << type << " thread_max = -INFINITY;" << std::endl;
    ks << "for (int i = threadIdx.x; i < row_size; i += blockDim.x) {" << std::endl;
    ks.setIndent(ks.indent() + 4);
    ks << "thread_max = fmax" << (type == "float" ? "f" : "") << "(thread_max, row_in[i]);" << std::endl;
    ks.setIndent(ks.indent() - 4);
    ks << "}" << std::endl;
    ks << std::endl;

    // Warp-level max reduction
    ks << "// Warp-level max reduction" << std::endl;
    ks << "for (int mask = 16; mask > 0; mask >>= 1) {" << std::endl;
    ks.setIndent(ks.indent() + 4);
    ks << "thread_max = fmax" << (type == "float" ? "f" : "")
       << "(thread_max, __shfl_xor_sync(0xFFFFFFFF, thread_max, mask));" << std::endl;
    ks.setIndent(ks.indent() - 4);
    ks << "}" << std::endl;
    ks << std::endl;

    // Cross-warp max reduction
    ks << "// Cross-warp max reduction" << std::endl;
    ks << "if (lane_id == 0) sdata[warp_id] = thread_max;" << std::endl;
    ks << "__syncthreads();" << std::endl;
    ks << type << " row_max = (threadIdx.x < num_warps) ? sdata[threadIdx.x] : (" << type << ")(-INFINITY);"
       << std::endl;
    ks << "for (int mask = 16; mask > 0; mask >>= 1) {" << std::endl;
    ks.setIndent(ks.indent() + 4);
    ks << "row_max = fmax" << (type == "float" ? "f" : "") << "(row_max, __shfl_xor_sync(0xFFFFFFFF, row_max, mask));"
       << std::endl;
    ks.setIndent(ks.indent() - 4);
    ks << "}" << std::endl;
    ks << "if (threadIdx.x == 0) sdata[0] = row_max;" << std::endl;
    ks << "__syncthreads();" << std::endl;
    ks << "row_max = sdata[0];" << std::endl;
    ks << std::endl;

    // Phase 2: exp and sum
    ks << "// Phase 2: exp(x - max) and sum" << std::endl;
    ks << type << " thread_sum = 0;" << std::endl;
    ks << "for (int i = threadIdx.x; i < row_size; i += blockDim.x) {" << std::endl;
    ks.setIndent(ks.indent() + 4);
    ks << type << " val = exp" << (type == "float" ? "f" : "") << "(row_in[i] - row_max);" << std::endl;
    ks << "row_out[i] = val;" << std::endl;
    ks << "thread_sum += val;" << std::endl;
    ks.setIndent(ks.indent() - 4);
    ks << "}" << std::endl;
    ks << std::endl;

    // Warp-level sum reduction
    ks << "// Warp-level sum reduction" << std::endl;
    ks << "for (int mask = 16; mask > 0; mask >>= 1) {" << std::endl;
    ks.setIndent(ks.indent() + 4);
    ks << "thread_sum += __shfl_xor_sync(0xFFFFFFFF, thread_sum, mask);" << std::endl;
    ks.setIndent(ks.indent() - 4);
    ks << "}" << std::endl;
    ks << std::endl;

    // Cross-warp sum reduction
    ks << "// Cross-warp sum reduction" << std::endl;
    ks << "if (lane_id == 0) sdata[warp_id] = thread_sum;" << std::endl;
    ks << "__syncthreads();" << std::endl;
    ks << type << " row_sum = (threadIdx.x < num_warps) ? sdata[threadIdx.x] : 0;" << std::endl;
    ks << "for (int mask = 16; mask > 0; mask >>= 1) {" << std::endl;
    ks.setIndent(ks.indent() + 4);
    ks << "row_sum += __shfl_xor_sync(0xFFFFFFFF, row_sum, mask);" << std::endl;
    ks.setIndent(ks.indent() - 4);
    ks << "}" << std::endl;
    ks << "if (threadIdx.x == 0) sdata[0] = row_sum;" << std::endl;
    ks << "__syncthreads();" << std::endl;
    ks << "row_sum = sdata[0];" << std::endl;
    ks << std::endl;

    // Phase 3: normalize
    ks << "// Phase 3: normalize" << std::endl;
    ks << "for (int i = threadIdx.x; i < row_size; i += blockDim.x) {" << std::endl;
    ks.setIndent(ks.indent() + 4);
    ks << "row_out[i] /= row_sum;" << std::endl;
    ks.setIndent(ks.indent() - 4);
    ks << "}" << std::endl;

    ks.setIndent(ks.indent() - 4);
    ks << "}" << std::endl;
}

// Strided softmax kernel: reduced axis is NOT innermost, so a softmax group's elements are
// `inner` apart in memory (layout [outer, reduce, inner]). Each thread owns one entire softmax
// group (one (outer, inner_idx) column). Consecutive threads map to consecutive `inner_idx`
// values, so for a fixed reduction index r all lanes of a warp touch consecutive addresses
// (outer*row_size*inner + r*inner + inner_idx) => fully coalesced global-memory accesses.
static void
emit_softmax_kernel_strided(codegen::PrettyPrinter& ks, const std::string& kernel_name, const std::string& type) {
    std::string fsuf = (type == "float" ? "f" : "");

    ks << "__global__ void " << kernel_name << "(const " << type << "* __restrict__ input, " << type
       << "* __restrict__ output, int num_groups, int row_size, int inner) {" << std::endl;
    ks.setIndent(ks.indent() + 4);

    // Grid-stride loop over softmax groups so any grid size is valid.
    ks << "for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < num_groups; g += gridDim.x * blockDim.x) {"
       << std::endl;
    ks.setIndent(ks.indent() + 4);

    ks << "int __outer = g / inner;" << std::endl;
    ks << "int __inner_idx = g % inner;" << std::endl;
    ks << "const " << type << "* col_in = input + (size_t)__outer * row_size * inner + __inner_idx;" << std::endl;
    ks << type << "* col_out = output + (size_t)__outer * row_size * inner + __inner_idx;" << std::endl;
    ks << std::endl;

    // Phase 1: max over the reduced axis (coalesced across the warp for each r)
    ks << type << " m = -INFINITY;" << std::endl;
    ks << "for (int r = 0; r < row_size; ++r) {" << std::endl;
    ks.setIndent(ks.indent() + 4);
    ks << "m = fmax" << fsuf << "(m, col_in[(size_t)r * inner]);" << std::endl;
    ks.setIndent(ks.indent() - 4);
    ks << "}" << std::endl;
    ks << std::endl;

    // Phase 2: exp(x - max), write to output, accumulate sum
    ks << type << " s = 0;" << std::endl;
    ks << "for (int r = 0; r < row_size; ++r) {" << std::endl;
    ks.setIndent(ks.indent() + 4);
    ks << type << " v = exp" << fsuf << "(col_in[(size_t)r * inner] - m);" << std::endl;
    ks << "col_out[(size_t)r * inner] = v;" << std::endl;
    ks << "s += v;" << std::endl;
    ks.setIndent(ks.indent() - 4);
    ks << "}" << std::endl;
    ks << std::endl;

    // Phase 3: normalize
    ks << type << " inv = (" << type << ")1 / s;" << std::endl;
    ks << "for (int r = 0; r < row_size; ++r) {" << std::endl;
    ks.setIndent(ks.indent() + 4);
    ks << "col_out[(size_t)r * inner] *= inv;" << std::endl;
    ks.setIndent(ks.indent() - 4);
    ks << "}" << std::endl;

    ks.setIndent(ks.indent() - 4);
    ks << "}" << std::endl; // grid-stride loop

    ks.setIndent(ks.indent() - 4);
    ks << "}" << std::endl;
}

static void compute_row_dims(
    const sdfg::math::tensor::SoftmaxNode& node,
    codegen::LanguageExtension& lang,
    std::string& num_rows_str,
    std::string& row_size_str,
    std::string& inner_str
) {
    auto& shape = node.shape();
    auto& axes = node.axes();
    int64_t ndim = static_cast<int64_t>(shape.size());

    // Normalize axes to positive
    std::set<int64_t> reduce_axes;
    for (auto a : axes) {
        reduce_axes.insert(a < 0 ? a + ndim : a);
    }

    // Decompose the (row-major) tensor into (outer, reduce, inner):
    //   - outer:  product of dims before the first reduced axis
    //   - reduce: product of dims spanning the reduced axes (row_size of each softmax group)
    //   - inner:  product of dims after the last reduced axis (memory stride between
    //             consecutive reduced elements)
    // When the reduced axes are trailing (inner == 1) the softmax groups are contiguous in
    // memory; otherwise they are strided by `inner`.
    int64_t reduce_min = ndim;
    int64_t reduce_max = -1;
    for (auto a : reduce_axes) {
        reduce_min = std::min(reduce_min, a);
        reduce_max = std::max(reduce_max, a);
    }

    symbolic::Expression outer = symbolic::one();
    symbolic::Expression reduce = symbolic::one();
    symbolic::Expression inner = symbolic::one();
    for (int64_t i = 0; i < ndim; ++i) {
        if (i < reduce_min) {
            outer = symbolic::mul(outer, shape[i]);
        } else if (i > reduce_max) {
            inner = symbolic::mul(inner, shape[i]);
        } else {
            reduce = symbolic::mul(reduce, shape[i]);
        }
    }

    // num_rows = number of independent softmax groups = outer * inner
    num_rows_str = lang.expression(symbolic::mul(outer, inner));
    row_size_str = lang.expression(reduce);
    inner_str = lang.expression(inner);
}

static std::string get_type_string(types::PrimitiveType prim_type) {
    switch (prim_type) {
        case types::PrimitiveType::Float:
            return "float";
        case types::PrimitiveType::Double:
            return "double";
        default:
            throw std::runtime_error("Unsupported primitive type for CUDA softmax dispatcher");
    }
}

static void dispatch_softmax_common(
    codegen::CodegenOutput& out,
    std::vector<codegen::DispatchInput>& inputs,
    codegen::LanguageExtension& language_extension,
    const sdfg::math::tensor::SoftmaxNode& node,
    const data_flow::DataFlowGraph& data_flow_graph,
    const std::string& input_ptr,
    const std::string& output_ptr
) {
    auto prim_type = node.primitive_type(data_flow_graph);
    std::string type = get_type_string(prim_type);

    std::string num_rows_str, row_size_str, inner_str;
    compute_row_dims(node, language_extension, num_rows_str, row_size_str, inner_str);

    // When the reduced axes are trailing, the softmax groups are contiguous (inner == 1) and we
    // can use the fast contiguous kernel exclusively. Otherwise we also need the strided kernel.
    bool inner_is_one = (inner_str == "1");

    std::string kernel_name = "softmax_kernel_" + std::to_string(node.element_id());
    std::string kernel_name_strided = kernel_name + "_strided";

    out.library_snippet_factory.add_global("#include <cuda.h>");
    out.library_snippet_factory.add_global("#include <math.h>");

    // Forward-declare kernel(s) in globals
    out.globals_stream << "__global__ void " << kernel_name << "(const " << type << "* __restrict__ input, " << type
                       << "* __restrict__ output, int num_rows, int row_size);" << std::endl;
    if (!inner_is_one) {
        out.globals_stream << "__global__ void " << kernel_name_strided << "(const " << type << "* __restrict__ input, "
                           << type << "* __restrict__ output, int num_rows, int row_size, int inner);" << std::endl;
    }

    // Emit kernel(s) to .cu file
    auto& kernel_stream = out.library_snippet_factory.require(kernel_name, "cu", true).stream();
    kernel_stream << "#include " << out.library_snippet_factory.header_path().filename() << std::endl << std::endl;
    emit_softmax_kernel_contiguous(kernel_stream, kernel_name, type);
    if (!inner_is_one) {
        kernel_stream << std::endl;
        emit_softmax_kernel_strided(kernel_stream, kernel_name_strided, type);
    }

    // Emit kernel call
    out.stream << "{" << std::endl;
    out.stream.setIndent(out.stream.indent() + 4);

    out.stream << "int __softmax_num_rows = (int)(" << num_rows_str << ");" << std::endl;
    out.stream << "int __softmax_row_size = (int)(" << row_size_str << ");" << std::endl;
    if (!inner_is_one) {
        out.stream << "int __softmax_inner = (int)(" << inner_str << ");" << std::endl;
    }

    // Launch config for the contiguous kernel: one block per row, block sized to the row.
    out.stream << "int __softmax_block_size = " << SOFTMAX_BLOCK_SIZE << ";" << std::endl;
    out.stream << "if (__softmax_row_size < __softmax_block_size) __softmax_block_size = __softmax_row_size;"
               << std::endl;
    // Round up to multiple of 32 (warp size)
    out.stream << "__softmax_block_size = ((__softmax_block_size + 31) / 32) * 32;" << std::endl;
    out.stream << "int __softmax_num_warps = __softmax_block_size / 32;" << std::endl;
    out.stream << "size_t __softmax_smem = __softmax_num_warps * sizeof(" << type << ");" << std::endl;

    auto emit_contiguous_launch = [&]() {
        out.stream << kernel_name << "<<<__softmax_num_rows, __softmax_block_size, __softmax_smem>>>(" << input_ptr
                   << ", " << output_ptr << ", __softmax_num_rows, __softmax_row_size);" << std::endl;
    };

    if (inner_is_one) {
        emit_contiguous_launch();
    } else {
        // Runtime dispatch: prefer the fast contiguous kernel whenever the stride collapses to 1.
        out.stream << "if (__softmax_inner == 1) {" << std::endl;
        out.stream.setIndent(out.stream.indent() + 4);
        emit_contiguous_launch();
        out.stream.setIndent(out.stream.indent() - 4);
        out.stream << "} else {" << std::endl;
        out.stream.setIndent(out.stream.indent() + 4);
        // Coalesced strided kernel: one thread per softmax group, grid-stride over groups.
        out.stream << "int __softmax_strided_block = " << SOFTMAX_BLOCK_SIZE << ";" << std::endl;
        out.stream << "int __softmax_strided_grid = (__softmax_num_rows + __softmax_strided_block - 1) / "
                      "__softmax_strided_block;"
                   << std::endl;
        out.stream << kernel_name_strided << "<<<__softmax_strided_grid, __softmax_strided_block>>>(" << input_ptr
                   << ", " << output_ptr << ", __softmax_num_rows, __softmax_row_size, __softmax_inner);" << std::endl;
        out.stream.setIndent(out.stream.indent() - 4);
        out.stream << "}" << std::endl;
    }

    check_cuda_kernel_launch_errors(out.stream, language_extension, false);

    out.stream.setIndent(out.stream.indent() - 4);
    out.stream << "}" << std::endl;
}

// WithTransfers

SoftmaxNodeDispatcher_CUDAWithTransfers::SoftmaxNodeDispatcher_CUDAWithTransfers(
    codegen::LanguageExtension& language_extension,
    const Function& function,
    const data_flow::DataFlowGraph& data_flow_graph,
    const sdfg::math::tensor::SoftmaxNode& node
)
    : codegen::LibraryNodeDispatcher(language_extension, function, data_flow_graph, node) {}

void SoftmaxNodeDispatcher_CUDAWithTransfers::dispatch_code_with_edges(
    codegen::CodegenOutput& out,
    std::vector<codegen::DispatchInput>& inputs,
    std::vector<codegen::DispatchOutput>& outputs
) {
    auto& node = static_cast<const sdfg::math::tensor::SoftmaxNode&>(this->node_);
    auto prim_type = node.primitive_type(this->data_flow_graph_);
    std::string type = get_type_string(prim_type);

    // Connectors: inputs_={"Y", "X"} → inputs[0]=Y (output buffer), inputs[1]=X (input data)
    auto& y_expr = inputs.at(0).expr;
    auto& x_expr = inputs.at(1).expr;

    std::string num_rows_str, row_size_str, inner_str;
    compute_row_dims(node, this->language_extension_, num_rows_str, row_size_str, inner_str);

    std::string total_size = "((size_t)(" + num_rows_str + ") * (size_t)(" + row_size_str + ")) * sizeof(" + type + ")";

    out.stream << "{" << std::endl;
    out.stream.setIndent(out.stream.indent() + 4);

    out.stream << "cudaError_t err_cuda;" << std::endl;
    out.stream << type << " *d_input, *d_output;" << std::endl;
    out.stream << "size_t __softmax_total_bytes = " << total_size << ";" << std::endl;

    out.stream << "err_cuda = cudaMalloc((void**) &d_input, __softmax_total_bytes);" << std::endl;
    cuda_error_checking(out.stream, this->language_extension_, "err_cuda");
    out.stream << "err_cuda = cudaMalloc((void**) &d_output, __softmax_total_bytes);" << std::endl;
    cuda_error_checking(out.stream, this->language_extension_, "err_cuda");

    out.stream << "err_cuda = cudaMemcpy(d_input, " << x_expr << ", __softmax_total_bytes, cudaMemcpyHostToDevice);"
               << std::endl;
    cuda_error_checking(out.stream, this->language_extension_, "err_cuda");

    dispatch_softmax_common(out, inputs, this->language_extension_, node, this->data_flow_graph_, "d_input", "d_output");

    out.stream << "err_cuda = cudaMemcpy(" << y_expr << ", d_output, __softmax_total_bytes, cudaMemcpyDeviceToHost);"
               << std::endl;
    cuda_error_checking(out.stream, this->language_extension_, "err_cuda");

    out.stream << "err_cuda = cudaFree(d_input);" << std::endl;
    cuda_error_checking(out.stream, this->language_extension_, "err_cuda");
    out.stream << "err_cuda = cudaFree(d_output);" << std::endl;
    cuda_error_checking(out.stream, this->language_extension_, "err_cuda");

    out.stream.setIndent(out.stream.indent() - 4);
    out.stream << "}" << std::endl;
}

// WithoutTransfers

SoftmaxNodeDispatcher_CUDAWithoutTransfers::SoftmaxNodeDispatcher_CUDAWithoutTransfers(
    codegen::LanguageExtension& language_extension,
    const Function& function,
    const data_flow::DataFlowGraph& data_flow_graph,
    const sdfg::math::tensor::SoftmaxNode& node
)
    : codegen::LibraryNodeDispatcher(language_extension, function, data_flow_graph, node) {}

void SoftmaxNodeDispatcher_CUDAWithoutTransfers::dispatch_code_with_edges(
    codegen::CodegenOutput& out,
    std::vector<codegen::DispatchInput>& inputs,
    std::vector<codegen::DispatchOutput>& outputs
) {
    auto& node = static_cast<const sdfg::math::tensor::SoftmaxNode&>(this->node_);

    // Connectors: inputs_={"Y", "X"} → inputs[0]=Y (output buffer), inputs[1]=X (input data)
    auto& y_expr = inputs.at(0).expr;
    auto& x_expr = inputs.at(1).expr;

    dispatch_softmax_common(out, inputs, this->language_extension_, node, this->data_flow_graph_, x_expr, y_expr);
}

} // namespace sdfg::cuda::tensor

1	#include "sdfg/targets/cuda/math/tensor/softmax.h"
2	#include "sdfg/symbolic/symbolic.h"
3	#include "sdfg/targets/cuda/cuda.h"
4
5	#include <algorithm>
6
7	namespace sdfg::cuda::tensor {
8
9	static constexpr int SOFTMAX_BLOCK_SIZE = 256;
10
11	// Contiguous softmax kernel: one block per row, warp-shuffle + shared-memory reductions.
12	// Used when the reduced axis is innermost (inner == 1), so a row's elements are contiguous
13	// in memory and threads within a warp access consecutive addresses (coalesced).
14	static void
15	emit_softmax_kernel_contiguous(codegen::PrettyPrinter& ks, const std::string& kernel_name, const std::string& type) {	5✔
16	ks << "__global__ void " << kernel_name << "(const " << type << "* __restrict__ input, " << type	5✔
17	<< "* __restrict__ output, int num_rows, int row_size) {" << std::endl;	5✔
18	ks.setIndent(ks.indent() + 4);	5✔
19
20	ks << "int row = blockIdx.x;" << std::endl;	5✔
21	ks << "if (row >= num_rows) return;" << std::endl;	5✔
22	ks << std::endl;	5✔
23	ks << "const " << type << "* row_in = input + (size_t)row * row_size;" << std::endl;	5✔
24	ks << type << "* row_out = output + (size_t)row * row_size;" << std::endl;	5✔
25	ks << std::endl;	5✔
26
27	// Shared memory for cross-warp reduction
28	ks << "extern __shared__ " << type << " sdata[];" << std::endl;	5✔
29	ks << "int lane_id = threadIdx.x & 31;" << std::endl;	5✔
30	ks << "int warp_id = threadIdx.x >> 5;" << std::endl;	5✔
31	ks << "int num_warps = (blockDim.x + 31) >> 5;" << std::endl;	5✔
32	ks << std::endl;	5✔
33
34	// Phase 1: find row max
35	ks << "// Phase 1: row max" << std::endl;	5✔
36	ks << type << " thread_max = -INFINITY;" << std::endl;	5✔
37	ks << "for (int i = threadIdx.x; i < row_size; i += blockDim.x) {" << std::endl;	5✔
38	ks.setIndent(ks.indent() + 4);	5✔
39	ks << "thread_max = fmax" << (type == "float" ? "f" : "") << "(thread_max, row_in[i]);" << std::endl;	5✔
40	ks.setIndent(ks.indent() - 4);	5✔
41	ks << "}" << std::endl;	5✔
42	ks << std::endl;	5✔
43
44	// Warp-level max reduction
45	ks << "// Warp-level max reduction" << std::endl;	5✔
46	ks << "for (int mask = 16; mask > 0; mask >>= 1) {" << std::endl;	5✔
47	ks.setIndent(ks.indent() + 4);	5✔
48	ks << "thread_max = fmax" << (type == "float" ? "f" : "")	5✔
49	<< "(thread_max, __shfl_xor_sync(0xFFFFFFFF, thread_max, mask));" << std::endl;	5✔
50	ks.setIndent(ks.indent() - 4);	5✔
51	ks << "}" << std::endl;	5✔
52	ks << std::endl;	5✔
53
54	// Cross-warp max reduction
55	ks << "// Cross-warp max reduction" << std::endl;	5✔
56	ks << "if (lane_id == 0) sdata[warp_id] = thread_max;" << std::endl;	5✔
57	ks << "__syncthreads();" << std::endl;	5✔
58	ks << type << " row_max = (threadIdx.x < num_warps) ? sdata[threadIdx.x] : (" << type << ")(-INFINITY);"	5✔
59	<< std::endl;	5✔
60	ks << "for (int mask = 16; mask > 0; mask >>= 1) {" << std::endl;	5✔
61	ks.setIndent(ks.indent() + 4);	5✔
62	ks << "row_max = fmax" << (type == "float" ? "f" : "") << "(row_max, __shfl_xor_sync(0xFFFFFFFF, row_max, mask));"	5✔
63	<< std::endl;	5✔
64	ks.setIndent(ks.indent() - 4);	5✔
65	ks << "}" << std::endl;	5✔
66	ks << "if (threadIdx.x == 0) sdata[0] = row_max;" << std::endl;	5✔
67	ks << "__syncthreads();" << std::endl;	5✔
68	ks << "row_max = sdata[0];" << std::endl;	5✔
69	ks << std::endl;	5✔
70
71	// Phase 2: exp and sum
72	ks << "// Phase 2: exp(x - max) and sum" << std::endl;	5✔
73	ks << type << " thread_sum = 0;" << std::endl;	5✔
74	ks << "for (int i = threadIdx.x; i < row_size; i += blockDim.x) {" << std::endl;	5✔
75	ks.setIndent(ks.indent() + 4);	5✔
76	ks << type << " val = exp" << (type == "float" ? "f" : "") << "(row_in[i] - row_max);" << std::endl;	5✔
77	ks << "row_out[i] = val;" << std::endl;	5✔
78	ks << "thread_sum += val;" << std::endl;	5✔
79	ks.setIndent(ks.indent() - 4);	5✔
80	ks << "}" << std::endl;	5✔
81	ks << std::endl;	5✔
82
83	// Warp-level sum reduction
84	ks << "// Warp-level sum reduction" << std::endl;	5✔
85	ks << "for (int mask = 16; mask > 0; mask >>= 1) {" << std::endl;	5✔
86	ks.setIndent(ks.indent() + 4);	5✔
87	ks << "thread_sum += __shfl_xor_sync(0xFFFFFFFF, thread_sum, mask);" << std::endl;	5✔
88	ks.setIndent(ks.indent() - 4);	5✔
89	ks << "}" << std::endl;	5✔
90	ks << std::endl;	5✔
91
92	// Cross-warp sum reduction
93	ks << "// Cross-warp sum reduction" << std::endl;	5✔
94	ks << "if (lane_id == 0) sdata[warp_id] = thread_sum;" << std::endl;	5✔
95	ks << "__syncthreads();" << std::endl;	5✔
96	ks << type << " row_sum = (threadIdx.x < num_warps) ? sdata[threadIdx.x] : 0;" << std::endl;	5✔
97	ks << "for (int mask = 16; mask > 0; mask >>= 1) {" << std::endl;	5✔
98	ks.setIndent(ks.indent() + 4);	5✔
99	ks << "row_sum += __shfl_xor_sync(0xFFFFFFFF, row_sum, mask);" << std::endl;	5✔
100	ks.setIndent(ks.indent() - 4);	5✔
101	ks << "}" << std::endl;	5✔
102	ks << "if (threadIdx.x == 0) sdata[0] = row_sum;" << std::endl;	5✔
103	ks << "__syncthreads();" << std::endl;	5✔
104	ks << "row_sum = sdata[0];" << std::endl;	5✔
105	ks << std::endl;	5✔
106
107	// Phase 3: normalize
108	ks << "// Phase 3: normalize" << std::endl;	5✔
109	ks << "for (int i = threadIdx.x; i < row_size; i += blockDim.x) {" << std::endl;	5✔
110	ks.setIndent(ks.indent() + 4);	5✔
111	ks << "row_out[i] /= row_sum;" << std::endl;	5✔
112	ks.setIndent(ks.indent() - 4);	5✔
113	ks << "}" << std::endl;	5✔
114
115	ks.setIndent(ks.indent() - 4);	5✔
116	ks << "}" << std::endl;	5✔
117	}	5✔
118
119	// Strided softmax kernel: reduced axis is NOT innermost, so a softmax group's elements are
120	// `inner` apart in memory (layout [outer, reduce, inner]). Each thread owns one entire softmax
121	// group (one (outer, inner_idx) column). Consecutive threads map to consecutive `inner_idx`
122	// values, so for a fixed reduction index r all lanes of a warp touch consecutive addresses
123	// (outerrow_sizeinner + r*inner + inner_idx) => fully coalesced global-memory accesses.
124	static void
NEW 125	emit_softmax_kernel_strided(codegen::PrettyPrinter& ks, const std::string& kernel_name, const std::string& type) {	×
NEW 126	std::string fsuf = (type == "float" ? "f" : "");	×
127
NEW 128	ks << "__global__ void " << kernel_name << "(const " << type << "* __restrict__ input, " << type	×
NEW 129	<< "* __restrict__ output, int num_groups, int row_size, int inner) {" << std::endl;	×
NEW 130	ks.setIndent(ks.indent() + 4);	×
131
132	// Grid-stride loop over softmax groups so any grid size is valid.
NEW 133	ks << "for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < num_groups; g += gridDim.x * blockDim.x) {"	×
NEW 134	<< std::endl;	×
NEW 135	ks.setIndent(ks.indent() + 4);	×
136
NEW 137	ks << "int __outer = g / inner;" << std::endl;	×
NEW 138	ks << "int __inner_idx = g % inner;" << std::endl;	×
NEW 139	ks << "const " << type << "* col_in = input + (size_t)__outer * row_size * inner + __inner_idx;" << std::endl;	×
NEW 140	ks << type << "* col_out = output + (size_t)__outer * row_size * inner + __inner_idx;" << std::endl;	×
NEW 141	ks << std::endl;	×
142
143	// Phase 1: max over the reduced axis (coalesced across the warp for each r)
NEW 144	ks << type << " m = -INFINITY;" << std::endl;	×
NEW 145	ks << "for (int r = 0; r < row_size; ++r) {" << std::endl;	×
NEW 146	ks.setIndent(ks.indent() + 4);	×
NEW 147	ks << "m = fmax" << fsuf << "(m, col_in[(size_t)r * inner]);" << std::endl;	×
NEW 148	ks.setIndent(ks.indent() - 4);	×
NEW 149	ks << "}" << std::endl;	×
NEW 150	ks << std::endl;	×
151
152	// Phase 2: exp(x - max), write to output, accumulate sum
NEW 153	ks << type << " s = 0;" << std::endl;	×
NEW 154	ks << "for (int r = 0; r < row_size; ++r) {" << std::endl;	×
NEW 155	ks.setIndent(ks.indent() + 4);	×
NEW 156	ks << type << " v = exp" << fsuf << "(col_in[(size_t)r * inner] - m);" << std::endl;	×
NEW 157	ks << "col_out[(size_t)r * inner] = v;" << std::endl;	×
NEW 158	ks << "s += v;" << std::endl;	×
NEW 159	ks.setIndent(ks.indent() - 4);	×
NEW 160	ks << "}" << std::endl;	×
NEW 161	ks << std::endl;	×
162
163	// Phase 3: normalize
NEW 164	ks << type << " inv = (" << type << ")1 / s;" << std::endl;	×
NEW 165	ks << "for (int r = 0; r < row_size; ++r) {" << std::endl;	×
NEW 166	ks.setIndent(ks.indent() + 4);	×
NEW 167	ks << "col_out[(size_t)r * inner] *= inv;" << std::endl;	×
NEW 168	ks.setIndent(ks.indent() - 4);	×
NEW 169	ks << "}" << std::endl;	×
170
NEW 171	ks.setIndent(ks.indent() - 4);	×
NEW 172	ks << "}" << std::endl; // grid-stride loop	×
173
NEW 174	ks.setIndent(ks.indent() - 4);	×
NEW 175	ks << "}" << std::endl;	×
NEW 176	}	×
177
178	static void compute_row_dims(
179	const sdfg::math::tensor::SoftmaxNode& node,
180	codegen::LanguageExtension& lang,
181	std::string& num_rows_str,
182	std::string& row_size_str,
183	std::string& inner_str
184	) {	7✔
185	auto& shape = node.shape();	7✔
186	auto& axes = node.axes();	7✔
187	int64_t ndim = static_cast<int64_t>(shape.size());	7✔
188
189	// Normalize axes to positive
190	std::set<int64_t> reduce_axes;	7✔
191	for (auto a : axes) {	7✔
192	reduce_axes.insert(a < 0 ? a + ndim : a);	7✔
193	}	7✔
194
195	// Decompose the (row-major) tensor into (outer, reduce, inner):
196	// - outer: product of dims before the first reduced axis
197	// - reduce: product of dims spanning the reduced axes (row_size of each softmax group)
198	// - inner: product of dims after the last reduced axis (memory stride between
199	// consecutive reduced elements)
200	// When the reduced axes are trailing (inner == 1) the softmax groups are contiguous in
201	// memory; otherwise they are strided by `inner`.
202	int64_t reduce_min = ndim;	7✔
203	int64_t reduce_max = -1;	7✔
204	for (auto a : reduce_axes) {	7✔
205	reduce_min = std::min(reduce_min, a);	7✔
206	reduce_max = std::max(reduce_max, a);	7✔
207	}	7✔
208
209	symbolic::Expression outer = symbolic::one();	7✔
210	symbolic::Expression reduce = symbolic::one();	7✔
211	symbolic::Expression inner = symbolic::one();	7✔
212	for (int64_t i = 0; i < ndim; ++i) {	27✔
213	if (i < reduce_min) {	20✔
214	outer = symbolic::mul(outer, shape[i]);	13✔
215	} else if (i > reduce_max) {	13✔
NEW 216	inner = symbolic::mul(inner, shape[i]);	×
217	} else {	7✔
218	reduce = symbolic::mul(reduce, shape[i]);	7✔
219	}	7✔
220	}	20✔
221
222	// num_rows = number of independent softmax groups = outer * inner
223	num_rows_str = lang.expression(symbolic::mul(outer, inner));	7✔
224	row_size_str = lang.expression(reduce);	7✔
225	inner_str = lang.expression(inner);	7✔
226	}	7✔
227
228	static std::string get_type_string(types::PrimitiveType prim_type) {	7✔
229	switch (prim_type) {	7✔
230	case types::PrimitiveType::Float:	7✔
231	return "float";	7✔
232	case types::PrimitiveType::Double:	×
233	return "double";	×
234	default:	×
235	throw std::runtime_error("Unsupported primitive type for CUDA softmax dispatcher");	×
236	}	7✔
237	}	7✔
238
239	static void dispatch_softmax_common(
240	codegen::CodegenOutput& out,
241	std::vector<codegen::DispatchInput>& inputs,
242	codegen::LanguageExtension& language_extension,
243	const sdfg::math::tensor::SoftmaxNode& node,
244	const data_flow::DataFlowGraph& data_flow_graph,
245	const std::string& input_ptr,
246	const std::string& output_ptr
247	) {	5✔
248	auto prim_type = node.primitive_type(data_flow_graph);	5✔
249	std::string type = get_type_string(prim_type);	5✔
250
251	std::string num_rows_str, row_size_str, inner_str;	5✔
252	compute_row_dims(node, language_extension, num_rows_str, row_size_str, inner_str);	5✔
253
254	// When the reduced axes are trailing, the softmax groups are contiguous (inner == 1) and we
255	// can use the fast contiguous kernel exclusively. Otherwise we also need the strided kernel.
256	bool inner_is_one = (inner_str == "1");	5✔
257
258	std::string kernel_name = "softmax_kernel_" + std::to_string(node.element_id());	5✔
259	std::string kernel_name_strided = kernel_name + "_strided";	5✔
260
261	out.library_snippet_factory.add_global("#include <cuda.h>");	5✔
262	out.library_snippet_factory.add_global("#include <math.h>");	5✔
263
264	// Forward-declare kernel(s) in globals
265	out.globals_stream << "__global__ void " << kernel_name << "(const " << type << "* __restrict__ input, " << type	5✔
266	<< "* __restrict__ output, int num_rows, int row_size);" << std::endl;	5✔
267	if (!inner_is_one) {	5✔
NEW 268	out.globals_stream << "__global__ void " << kernel_name_strided << "(const " << type << "* __restrict__ input, "	×
NEW 269	<< type << "* __restrict__ output, int num_rows, int row_size, int inner);" << std::endl;	×
NEW 270	}	×
271
272	// Emit kernel(s) to .cu file
273	auto& kernel_stream = out.library_snippet_factory.require(kernel_name, "cu", true).stream();	5✔
274	kernel_stream << "#include " << out.library_snippet_factory.header_path().filename() << std::endl << std::endl;	5✔
275	emit_softmax_kernel_contiguous(kernel_stream, kernel_name, type);	5✔
276	if (!inner_is_one) {	5✔
NEW 277	kernel_stream << std::endl;	×
NEW 278	emit_softmax_kernel_strided(kernel_stream, kernel_name_strided, type);	×
NEW 279	}	×
280
281	// Emit kernel call
282	out.stream << "{" << std::endl;	5✔
283	out.stream.setIndent(out.stream.indent() + 4);	5✔
284
285	out.stream << "int __softmax_num_rows = (int)(" << num_rows_str << ");" << std::endl;	5✔
286	out.stream << "int __softmax_row_size = (int)(" << row_size_str << ");" << std::endl;	5✔
287	if (!inner_is_one) {	5✔
NEW 288	out.stream << "int __softmax_inner = (int)(" << inner_str << ");" << std::endl;	×
NEW 289	}	×
290
291	// Launch config for the contiguous kernel: one block per row, block sized to the row.
292	out.stream << "int __softmax_block_size = " << SOFTMAX_BLOCK_SIZE << ";" << std::endl;	5✔
293	out.stream << "if (__softmax_row_size < __softmax_block_size) __softmax_block_size = __softmax_row_size;"	5✔
294	<< std::endl;	5✔
295	// Round up to multiple of 32 (warp size)
296	out.stream << "__softmax_block_size = ((__softmax_block_size + 31) / 32) * 32;" << std::endl;	5✔
297	out.stream << "int __softmax_num_warps = __softmax_block_size / 32;" << std::endl;	5✔
298	out.stream << "size_t __softmax_smem = __softmax_num_warps * sizeof(" << type << ");" << std::endl;	5✔
299
300	auto emit_contiguous_launch = [&]() {	5✔
301	out.stream << kernel_name << "<<<__softmax_num_rows, __softmax_block_size, __softmax_smem>>>(" << input_ptr	5✔
302	<< ", " << output_ptr << ", __softmax_num_rows, __softmax_row_size);" << std::endl;	5✔
303	};	5✔
304
305	if (inner_is_one) {	5✔
306	emit_contiguous_launch();	5✔
307	} else {	5✔
308	// Runtime dispatch: prefer the fast contiguous kernel whenever the stride collapses to 1.
NEW 309	out.stream << "if (__softmax_inner == 1) {" << std::endl;	×
NEW 310	out.stream.setIndent(out.stream.indent() + 4);	×
NEW 311	emit_contiguous_launch();	×
NEW 312	out.stream.setIndent(out.stream.indent() - 4);	×
NEW 313	out.stream << "} else {" << std::endl;	×
NEW 314	out.stream.setIndent(out.stream.indent() + 4);	×
315	// Coalesced strided kernel: one thread per softmax group, grid-stride over groups.
NEW 316	out.stream << "int __softmax_strided_block = " << SOFTMAX_BLOCK_SIZE << ";" << std::endl;	×
NEW 317	out.stream << "int __softmax_strided_grid = (__softmax_num_rows + __softmax_strided_block - 1) / "	×
NEW 318	"__softmax_strided_block;"	×
NEW 319	<< std::endl;	×
NEW 320	out.stream << kernel_name_strided << "<<<__softmax_strided_grid, __softmax_strided_block>>>(" << input_ptr	×
NEW 321	<< ", " << output_ptr << ", __softmax_num_rows, __softmax_row_size, __softmax_inner);" << std::endl;	×
NEW 322	out.stream.setIndent(out.stream.indent() - 4);	×
NEW 323	out.stream << "}" << std::endl;	×
NEW 324	}	×
325
326	check_cuda_kernel_launch_errors(out.stream, language_extension, false);	5✔
327
328	out.stream.setIndent(out.stream.indent() - 4);	5✔
329	out.stream << "}" << std::endl;	5✔
330	}	5✔
331
332	// WithTransfers
333
334	SoftmaxNodeDispatcher_CUDAWithTransfers::SoftmaxNodeDispatcher_CUDAWithTransfers(
335	codegen::LanguageExtension& language_extension,
336	const Function& function,
337	const data_flow::DataFlowGraph& data_flow_graph,
338	const sdfg::math::tensor::SoftmaxNode& node
339	)
340	: codegen::LibraryNodeDispatcher(language_extension, function, data_flow_graph, node) {}	2✔
341
342	void SoftmaxNodeDispatcher_CUDAWithTransfers::dispatch_code_with_edges(
343	codegen::CodegenOutput& out,
344	std::vector<codegen::DispatchInput>& inputs,
345	std::vector<codegen::DispatchOutput>& outputs
346	) {	2✔
347	auto& node = static_cast<const sdfg::math::tensor::SoftmaxNode&>(this->node_);	2✔
348	auto prim_type = node.primitive_type(this->data_flow_graph_);	2✔
349	std::string type = get_type_string(prim_type);	2✔
350
351	// Connectors: inputs_={"Y", "X"} → inputs[0]=Y (output buffer), inputs[1]=X (input data)
352	auto& y_expr = inputs.at(0).expr;	2✔
353	auto& x_expr = inputs.at(1).expr;	2✔
354
355	std::string num_rows_str, row_size_str, inner_str;	2✔
356	compute_row_dims(node, this->language_extension_, num_rows_str, row_size_str, inner_str);	2✔
357
358	std::string total_size = "((size_t)(" + num_rows_str + ") * (size_t)(" + row_size_str + ")) * sizeof(" + type + ")";	2✔
359
360	out.stream << "{" << std::endl;	2✔
361	out.stream.setIndent(out.stream.indent() + 4);	2✔
362
363	out.stream << "cudaError_t err_cuda;" << std::endl;	2✔
364	out.stream << type << " d_input, d_output;" << std::endl;	2✔
365	out.stream << "size_t __softmax_total_bytes = " << total_size << ";" << std::endl;	2✔
366
367	out.stream << "err_cuda = cudaMalloc((void**) &d_input, __softmax_total_bytes);" << std::endl;	2✔
368	cuda_error_checking(out.stream, this->language_extension_, "err_cuda");	2✔
369	out.stream << "err_cuda = cudaMalloc((void**) &d_output, __softmax_total_bytes);" << std::endl;	2✔
370	cuda_error_checking(out.stream, this->language_extension_, "err_cuda");	2✔
371
372	out.stream << "err_cuda = cudaMemcpy(d_input, " << x_expr << ", __softmax_total_bytes, cudaMemcpyHostToDevice);"	2✔
373	<< std::endl;	2✔
374	cuda_error_checking(out.stream, this->language_extension_, "err_cuda");	2✔
375
376	dispatch_softmax_common(out, inputs, this->language_extension_, node, this->data_flow_graph_, "d_input", "d_output");	2✔
377
378	out.stream << "err_cuda = cudaMemcpy(" << y_expr << ", d_output, __softmax_total_bytes, cudaMemcpyDeviceToHost);"	2✔
379	<< std::endl;	2✔
380	cuda_error_checking(out.stream, this->language_extension_, "err_cuda");	2✔
381
382	out.stream << "err_cuda = cudaFree(d_input);" << std::endl;	2✔
383	cuda_error_checking(out.stream, this->language_extension_, "err_cuda");	2✔
384	out.stream << "err_cuda = cudaFree(d_output);" << std::endl;	2✔
385	cuda_error_checking(out.stream, this->language_extension_, "err_cuda");	2✔
386
387	out.stream.setIndent(out.stream.indent() - 4);	2✔
388	out.stream << "}" << std::endl;	2✔
389	}	2✔
390
391	// WithoutTransfers
392
393	SoftmaxNodeDispatcher_CUDAWithoutTransfers::SoftmaxNodeDispatcher_CUDAWithoutTransfers(
394	codegen::LanguageExtension& language_extension,
395	const Function& function,
396	const data_flow::DataFlowGraph& data_flow_graph,
397	const sdfg::math::tensor::SoftmaxNode& node
398	)
399	: codegen::LibraryNodeDispatcher(language_extension, function, data_flow_graph, node) {}	3✔
400
401	void SoftmaxNodeDispatcher_CUDAWithoutTransfers::dispatch_code_with_edges(
402	codegen::CodegenOutput& out,
403	std::vector<codegen::DispatchInput>& inputs,
404	std::vector<codegen::DispatchOutput>& outputs
405	) {	3✔
406	auto& node = static_cast<const sdfg::math::tensor::SoftmaxNode&>(this->node_);	3✔
407
408	// Connectors: inputs_={"Y", "X"} → inputs[0]=Y (output buffer), inputs[1]=X (input data)
409	auto& y_expr = inputs.at(0).expr;	3✔
410	auto& x_expr = inputs.at(1).expr;	3✔
411
412	dispatch_softmax_common(out, inputs, this->language_extension_, node, this->data_flow_graph_, x_expr, y_expr);	3✔
413	}	3✔
414
415	} // namespace sdfg::cuda::tensor

daisytuner / docc / 30338860029

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