diff --git a/docs/zero-element-graphs.md b/docs/zero-element-graphs.md new file mode 100644 index 000000000..d8ddf70ee --- /dev/null +++ b/docs/zero-element-graphs.md @@ -0,0 +1,63 @@ +# Zero-element (no-op) graphs + +Graphs whose tensors contain a dimension of size 0 — for example SDPA with a +batch size of 0 — are supported by the frontend as no-ops +(see [issue #101](https://github.com/NVIDIA/cudnn-frontend/issues/101)). + +The cuDNN backend does not accept tensor descriptors with 0-sized dimensions. +Instead of lowering such graphs, the frontend detects them during `validate()`: +if the graph references at least one zero-element tensor and **every** +non-virtual output tensor is zero-element, no output byte would ever be +written, so the graph is flagged as a zero-element no-op. For such graphs: + +- `validate()`, `build_operation_graph()`, `create_execution_plans()`, + `check_support()`, `build_plans()`, and `build()` all succeed without + touching the backend. +- `get_workspace_size()` reports 0. +- `execute()` returns success without launching any work. Pointers in the + variant pack are ignored (zero-element tensors may have null pointers). +- `populate_cuda_graph()` leaves the provided CUDA graph empty (an empty CUDA + graph is valid and instantiable), and `update_cuda_graph()` is a no-op. +- `Graph::is_zero_element_graph()` (C++) / `pygraph.is_zero_element_graph()` + (Python) report whether the graph was flagged. + +## Example + +```cpp +namespace fe = cudnn_frontend; +fe::graph::Graph graph; +graph.set_io_data_type(fe::DataType_t::HALF) + .set_intermediate_data_type(fe::DataType_t::FLOAT) + .set_compute_data_type(fe::DataType_t::FLOAT); + +int64_t b = 0, h = 4, s_q = 64, s_kv = 64, d = 64; // batch size 0 +auto Q = graph.tensor(fe::graph::Tensor_attributes() + .set_dim({b, h, s_q, d}) + .set_stride({h * s_q * d, s_q * d, d, 1})); +// ... K, V ... + +auto [O, Stats] = graph.sdpa(Q, K, V, sdpa_options); +O->set_output(true).set_dim({b, h, s_q, d}).set_stride({h * s_q * d, s_q * d, d, 1}); + +// The full pipeline succeeds; execute() is a no-op. +auto status = graph.build(handle, {fe::HeurMode_t::A}); +``` + +## Unsupported mixes + +A graph that mixes zero-element tensors with **non-zero-element** output +tensors is rejected at `validate()` with `GRAPH_NOT_SUPPORTED` and a clear +message. For example, a matmul contracting over a dimension of size 0 +(`[1, m, 0] x [1, 0, n] -> [1, m, n]`) would require zero-filling the output, +which cuDNN does not do. + +## Limitations + +- Plan serialization (`serialize(std::vector&)` / + `deserialize(handle, data)`) is not supported for zero-element no-op graphs, + as there is no execution plan to serialize. JSON graph (structure) + serialization is unaffected: a graph deserialized from JSON is re-detected as + a no-op during its `validate()`. +- Runtime shape overrides (`override_uids` / `override_shapes` / + `override_strides`) cannot be applied to a graph built as a zero-element + no-op; build the graph with non-zero shapes instead. diff --git a/include/cudnn_frontend/graph_interface.h b/include/cudnn_frontend/graph_interface.h index 8e92853aa..a6029476a 100644 --- a/include/cudnn_frontend/graph_interface.h +++ b/include/cudnn_frontend/graph_interface.h @@ -52,6 +52,12 @@ class Graph : public ICudnn, public INode { int64_t fe_workspace_size = 0; uint64_t graph_uid; + // Set during validate() when every non-virtual output tensor is zero-element + // (has a dimension of size 0, e.g. SDPA with batch size 0). Executing such a + // graph writes nothing, so backend lowering is skipped and all build/execute + // stages become no-ops. The backend does not accept 0-sized dimensions. + bool is_zero_element_graph_ = false; + std::unordered_set> deserialized_tensor_properties; std::unordered_map deserialized_pass_by_value; std::unordered_map>> deserialized_workspace_modifications; @@ -93,6 +99,54 @@ class Graph : public ICudnn, public INode { return {error_code_t::OK, ""}; } + // Detects zero-element (no-op) graphs. Called from validate(), after shape + // inference has filled in all tensor dims. + // A graph is a no-op iff it references at least one zero-element tensor and + // every non-virtual output tensor is zero-element: no output byte would ever + // be written, so all build/execute stages can safely be skipped. + // A graph that mixes zero-element tensors with non-zero-element outputs + // (e.g. a matmul whose contracted dimension is 0, whose output would need + // to be zero-filled) is rejected here with a clear message instead of + // failing later in the backend with a generic CUDNN_STATUS_BAD_PARAM. + error_t + detect_zero_element_graph() { + is_zero_element_graph_ = false; + + bool has_zero_element_tensor = false; + bool all_outputs_zero_element = true; + for (auto const &input : full_graph_inputs) { + if (input && input->get_volume() == 0) { + has_zero_element_tensor = true; + } + } + for (auto const &output : full_graph_outputs) { + if (output == nullptr || output->get_is_virtual()) { + continue; + } + if (output->get_volume() == 0) { + has_zero_element_tensor = true; + } else { + all_outputs_zero_element = false; + } + } + + if (has_zero_element_tensor == false) { + return {error_code_t::OK, ""}; + } + + RETURN_CUDNN_FRONTEND_ERROR_IF( + all_outputs_zero_element == false, + error_code_t::GRAPH_NOT_SUPPORTED, + "Graph mixes zero-element tensors (a dimension of size 0) with non-zero-element output tensors. " + "This is not supported: cuDNN cannot compute non-empty outputs from zero-element inputs."); + + is_zero_element_graph_ = true; + CUDNN_FE_LOG_LABEL_ENDL( + "INFO: All output tensors are zero-element. Graph is treated as a no-op: build stages are skipped, " + "workspace size is 0, and execute() launches no work."); + return {error_code_t::OK, ""}; + } + error_t log_tensors_to_dump_(cudnnHandle_t handle, std::unordered_map const &tensor_uid_to_pointer_map) const { @@ -646,6 +700,12 @@ class Graph : public ICudnn, public INode { std::unordered_map &uid_to_device_ptrs, void *workspace, cudaGraph_t cudnn_cuda_graph) { + // Zero-element no-op graph: populate_cuda_graph left the cuda graph empty; nothing to update. + if (is_zero_element_graph_) { + CUDNN_FE_LOG_LABEL_ENDL("INFO: Skipping cuda graph update of zero-element no-op graph."); + return {error_code_t::OK, ""}; + } + // Initializes this cudnn graph RETURN_CUDNN_FRONTEND_ERROR_IF( cudnn_cuda_graph == nullptr, error_code_t::INVALID_VALUE, "cudnn_cuda_graph should not be a nullptr"); @@ -788,6 +848,13 @@ class Graph : public ICudnn, public INode { std::unordered_map &uid_to_device_ptrs, void *workspace, cudaGraph_t cudnn_cuda_graph) { + // Zero-element no-op graph: leave the provided cuda graph empty (an empty + // cuda graph is valid and instantiable). + if (is_zero_element_graph_) { + CUDNN_FE_LOG_LABEL_ENDL("INFO: Skipping cuda graph population of zero-element no-op graph."); + return {error_code_t::OK, ""}; + } + // Check if the cuda graph is empty size_t numNodes = 0; _CUDNN_CHECK_CUDA_ERROR(detail::cuda_graph_get_nodes(cudnn_cuda_graph, nullptr, &numNodes)); @@ -936,6 +1003,9 @@ class Graph : public ICudnn, public INode { CHECK_CUDNN_FRONTEND_ERROR(output->validate()); } + // Now that all shapes are known, detect zero-element (no-op) graphs. + CHECK_CUDNN_FRONTEND_ERROR(detect_zero_element_graph()); + // Get all the pre assigned uids CHECK_CUDNN_FRONTEND_ERROR(get_pre_assigned_uids(used_uids)); // Clear state @@ -951,6 +1021,11 @@ class Graph : public ICudnn, public INode { build_operation_graph() { CUDNN_FE_LOG_BANNER(" BUILD OP GRAPH WITHOUT HANDLE "); + if (is_zero_element_graph_) { + CUDNN_FE_LOG_BANNER(" SKIPPED BUILD OP GRAPH (ZERO-ELEMENT NO-OP GRAPH) "); + return {error_code_t::OK, ""}; + } + if (device_properties == nullptr) { return {error_code_t::ATTRIBUTE_NOT_SET, "Device properties are not set."}; } @@ -962,6 +1037,11 @@ class Graph : public ICudnn, public INode { build_operation_graph(cudnnHandle_t handle) { CUDNN_FE_LOG_BANNER(" BUILD OP GRAPH "); + if (is_zero_element_graph_) { + CUDNN_FE_LOG_BANNER(" SKIPPED BUILD OP GRAPH (ZERO-ELEMENT NO-OP GRAPH) "); + return {error_code_t::OK, ""}; + } + CUDNN_FE_LOG_BANNER(" 1/4 INFER PROPERTIES OF NODES "); // expand composite nodes @@ -1061,6 +1141,13 @@ class Graph : public ICudnn, public INode { error_t get_workspace_size_plan_at_index(int64_t plan_index, int64_t &cudnn_workspace_size) const { + // Zero-element no-op graphs have no plans and need no workspace. + if (is_zero_element_graph_) { + cudnn_workspace_size = 0; + CUDNN_FE_LOG_LABEL_ENDL("INFO: get_workspace_size() is 0 (zero-element no-op graph)"); + return {error_code_t::OK, ""}; + } + // OSS SDPA engine workspace: 16 bytes for tile_id_counter if (plan_index == graph::Execution_plan_list::OSS_SDPA_ENGINE_CANDIDATE) { cudnn_workspace_size = fe_workspace_size + experimental::Sm90SdpaPrefillEngine::get_workspace_size(); @@ -1093,6 +1180,17 @@ class Graph : public ICudnn, public INode { std::vector const &override_uids, std::vector> const &override_shapes, std::vector> const &override_strides) const { + // Zero-element no-op graphs have no plans and need no workspace. + if (is_zero_element_graph_) { + RETURN_CUDNN_FRONTEND_ERROR_IF(!override_uids.empty(), + error_code_t::GRAPH_NOT_SUPPORTED, + "Graph was built as a zero-element no-op; it has no execution plan to " + "apply shape overrides to. Build the graph with non-zero shapes instead."); + cudnn_workspace_size = 0; + CUDNN_FE_LOG_LABEL_ENDL("INFO: get_workspace_size() is 0 (zero-element no-op graph)"); + return {error_code_t::OK, ""}; + } + RETURN_CUDNN_FRONTEND_ERROR_IF(override_uids.size() != override_shapes.size(), error_code_t::INVALID_VALUE, "override_uids and override_shapes must have the same size."); @@ -1218,6 +1316,11 @@ class Graph : public ICudnn, public INode { void *user_impl = nullptr) { (void)user_impl; // reserved for future use + if (is_zero_element_graph_) { + CUDNN_FE_LOG_LABEL_ENDL("INFO: Skipping autotune of zero-element no-op graph."); + return {error_code_t::OK, ""}; + } + const int maxIterCount = 100; const float threshold = 0.95f; @@ -1373,6 +1476,16 @@ class Graph : public ICudnn, public INode { std::vector const &override_uids = {}, std::vector> const &override_shapes = {}, std::vector> const &override_strides = {}) const { + // Zero-element no-op graph: all outputs are zero-element, nothing to compute. + if (is_zero_element_graph_) { + RETURN_CUDNN_FRONTEND_ERROR_IF(!override_uids.empty(), + error_code_t::GRAPH_NOT_SUPPORTED, + "Graph was built as a zero-element no-op; it has no execution plan to " + "apply shape overrides to. Build the graph with non-zero shapes instead."); + CUDNN_FE_LOG_LABEL_ENDL("INFO: Skipping execution of zero-element no-op graph."); + return {error_code_t::OK, ""}; + } + if (!varpack_prep_state->prepared.load(std::memory_order_acquire)) { CHECK_CUDNN_FRONTEND_ERROR(const_cast(this)->prepare_variant_pack_template()); } @@ -1404,6 +1517,16 @@ class Graph : public ICudnn, public INode { std::vector const &override_uids = {}, std::vector> const &override_shapes = {}, std::vector> const &override_strides = {}) const { + // Zero-element no-op graph: all outputs are zero-element, nothing to compute. + if (is_zero_element_graph_) { + RETURN_CUDNN_FRONTEND_ERROR_IF(!override_uids.empty(), + error_code_t::GRAPH_NOT_SUPPORTED, + "Graph was built as a zero-element no-op; it has no execution plan to " + "apply shape overrides to. Build the graph with non-zero shapes instead."); + CUDNN_FE_LOG_LABEL_ENDL("INFO: Skipping execution of zero-element no-op graph."); + return {error_code_t::OK, ""}; + } + // Lazy init: prepare template if not done (e.g. deserialized graphs, build_plan_at_index) if (!varpack_prep_state->prepared.load(std::memory_order_acquire)) { CHECK_CUDNN_FRONTEND_ERROR(const_cast(this)->prepare_variant_pack_template()); @@ -1502,6 +1625,11 @@ class Graph : public ICudnn, public INode { error_t warmup(cudnnHandle_t handle) { + if (is_zero_element_graph_) { + CUDNN_FE_LOG_LABEL_ENDL("INFO: Skipping warmup of zero-element no-op graph."); + return {error_code_t::OK, ""}; + } + cudaStream_t fake_stream; cudaStream_t original_stream; @@ -1610,6 +1738,9 @@ class Graph : public ICudnn, public INode { error_t serialize(std::vector &data, bool serialize_structure = true) const { CUDNN_FE_LOG_BANNER(" SERIALIZE PLAN "); + RETURN_CUDNN_FRONTEND_ERROR_IF(is_zero_element_graph_, + error_code_t::GRAPH_NOT_SUPPORTED, + "Zero-element no-op graphs have no execution plan to serialize."); #ifndef CUDNN_FRONTEND_SKIP_JSON_LIB json j; // Optionally serialize the graph structure (nodes/tensors). @@ -2119,9 +2250,24 @@ class Graph : public ICudnn, public INode { return check_support(); } + // Returns true if validate() determined this graph to be a zero-element no-op: + // it references zero-element tensors (a dimension of size 0, e.g. batch size 0) + // and every non-virtual output tensor is zero-element. Such graphs skip backend + // lowering, report a workspace size of 0, and execute() launches no work. + bool + is_zero_element_graph() const { + return is_zero_element_graph_; + } + // overload for deviceless AoT compilation error_t check_support() { + // Zero-element no-op graphs are always supported: no backend plan is needed. + if (is_zero_element_graph_) { + CUDNN_FE_LOG_LABEL_ENDL("INFO: check_support() OK for zero-element no-op graph."); + return {error_code_t::OK, ""}; + } + // Check OSS engine first if registered CHECK_CUDNN_FRONTEND_ERROR(context.populate_sm_version_from_device()); @@ -2826,6 +2972,11 @@ inline error_t Graph::create_execution_plans(std::vector const &mode) { CUDNN_FE_LOG_BANNER(" CREATE EXECUTION PLANS (HEURISTICS QUERY) "); + if (is_zero_element_graph_) { + CUDNN_FE_LOG_LABEL_ENDL("INFO: Skipping execution plan creation of zero-element no-op graph."); + return {error_code_t::OK, ""}; + } + // CHECK IF NEED TO OVERRIDE HEURISTICS QUERY for (auto &sub_node : sub_nodes) { if (auto [engine_id, user_knobs] = sub_node->override_heuristics_query(); engine_id != -1) { @@ -2896,6 +3047,12 @@ Graph::create_execution_plan(int64_t const engine_id, std::unordered_map(policy) << " "); #endif + if (is_zero_element_graph_) { + CUDNN_FE_LOG_BANNER(" SKIPPED BUILD PLANS (ZERO-ELEMENT NO-OP GRAPH) "); + return {error_code_t::OK, ""}; + } + // Build OSS SDPA engine if it passed check_support if (plans.has_oss_sdpa_engine()) { auto oss_status = plans.build_oss_sdpa_engine(); diff --git a/llms.txt b/llms.txt index 5439d4a8f..3e4d0b406 100644 --- a/llms.txt +++ b/llms.txt @@ -39,6 +39,7 @@ Published documentation: https://docs.nvidia.com/deeplearning/cudnn/latest/devel - [Dynamic kernel cache](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/dynamic-kernel-cache.md) - [Custom execution plans](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/custom-execution-plan.md) - [Compile-time constants](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/CompileTimeConstants.md) +- [Zero-element (no-op) graphs](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/zero-element-graphs.md) - [Adding PyTorch custom ops](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/adding_torch_custom_ops.md) - [Python graph and execution backends](https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/python_graph_and_execution_backends.md) diff --git a/python/pygraph/pygraph.cpp b/python/pygraph/pygraph.cpp index ed04f9b19..f4fc8b13e 100644 --- a/python/pygraph/pygraph.cpp +++ b/python/pygraph/pygraph.cpp @@ -544,6 +544,11 @@ PyGraph::check_support() { throw_if(status.is_bad(), status.get_code(), status.get_message()); } +bool +PyGraph::is_zero_element_graph() { + return graph->is_zero_element_graph(); +} + int64_t PyGraph::get_workspace_size() { int64_t workspace = 0; @@ -1280,6 +1285,14 @@ init_pygraph_submodule(py::module_& m) { engine_id (int): The ID of the engine to query knob configurations for. )pbdoc") .def("check_support", &PyGraph::check_support) + .def("is_zero_element_graph", + &PyGraph::is_zero_element_graph, + R"pbdoc( + Returns True if validate() determined this graph to be a zero-element no-op: + it references zero-element tensors (a dimension of size 0, e.g. batch size 0) + and every non-virtual output tensor is zero-element. Such graphs skip backend + lowering, report a workspace size of 0, and execute() launches no work. + )pbdoc") .def("build_plans", &PyGraph::build_plans, py::arg("policy") = cudnn_frontend::BuildPlanPolicy_t::HEURISTICS_CHOICE) diff --git a/python/pygraph/pygraph.h b/python/pygraph/pygraph.h index 4019f3560..6cb01ce17 100644 --- a/python/pygraph/pygraph.h +++ b/python/pygraph/pygraph.h @@ -672,6 +672,9 @@ class PyGraph { void check_support(); + bool + is_zero_element_graph(); + void build(std::vector const&); diff --git a/test/cpp/validate.cpp b/test/cpp/validate.cpp index 426f9392c..9f9b7644a 100644 --- a/test/cpp/validate.cpp +++ b/test/cpp/validate.cpp @@ -121,4 +121,98 @@ TEST_CASE("Multiple validation", "[graph][validate]") { REQUIRE(graph.validate().is_good()); REQUIRE(graph.validate().is_good()); +} + +TEST_CASE("Zero element graph is a no-op", "[graph][validate][zero_element]") { + namespace fe = cudnn_frontend; + fe::graph::Graph graph; + + graph.set_io_data_type(fe::DataType_t::HALF) + .set_intermediate_data_type(fe::DataType_t::FLOAT) + .set_compute_data_type(fe::DataType_t::FLOAT); + + auto A = graph.tensor(fe::graph::Tensor_attributes().set_name("A").set_dim({0, 8, 16}).set_stride({128, 16, 1})); + auto B = graph.tensor(fe::graph::Tensor_attributes().set_name("B").set_dim({0, 8, 16}).set_stride({128, 16, 1})); + + auto C = graph.pointwise(A, B, fe::graph::Pointwise_attributes().set_mode(fe::PointwiseMode_t::ADD)); + C->set_output(true); + + REQUIRE(graph.validate().is_good()); + REQUIRE(graph.is_zero_element_graph()); + + // The entire pipeline is a no-op and succeeds without a device or handle. + REQUIRE(graph.build_operation_graph(nullptr).is_good()); + REQUIRE(graph.create_execution_plans({fe::HeurMode_t::A}).is_good()); + REQUIRE(graph.check_support().is_good()); + REQUIRE(graph.build_plans().is_good()); + REQUIRE(graph.get_workspace_size() == 0); + + std::unordered_map, void *> variant_pack = { + {A, nullptr}, {B, nullptr}, {C, nullptr}}; + REQUIRE(graph.execute(nullptr, variant_pack, nullptr).is_good()); +} + +TEST_CASE("SDPA with batch size 0 is a no-op", "[graph][sdpa][validate][zero_element]") { + namespace fe = cudnn_frontend; + + int64_t const b = 0, h = 4, s_q = 64, s_kv = 64, d = 64; + + fe::graph::Graph graph; + graph.set_io_data_type(fe::DataType_t::HALF) + .set_intermediate_data_type(fe::DataType_t::FLOAT) + .set_compute_data_type(fe::DataType_t::FLOAT); + + auto Q = graph.tensor( + fe::graph::Tensor_attributes().set_name("Q").set_dim({b, h, s_q, d}).set_stride({h * s_q * d, s_q * d, d, 1})); + auto K = graph.tensor(fe::graph::Tensor_attributes() + .set_name("K") + .set_dim({b, h, s_kv, d}) + .set_stride({h * s_kv * d, s_kv * d, d, 1})); + auto V = graph.tensor(fe::graph::Tensor_attributes() + .set_name("V") + .set_dim({b, h, s_kv, d}) + .set_stride({h * s_kv * d, s_kv * d, d, 1})); + + auto sdpa_options = fe::graph::SDPA_attributes().set_name("sdpa").set_generate_stats(false).set_attn_scale(0.125f); + + auto [O, Stats] = graph.sdpa(Q, K, V, sdpa_options); + O->set_output(true).set_dim({b, h, s_q, d}).set_stride({h * s_q * d, s_q * d, d, 1}); + REQUIRE(Stats == nullptr); + + REQUIRE(graph.validate().is_good()); + REQUIRE(graph.is_zero_element_graph()); + + REQUIRE(graph.build_operation_graph(nullptr).is_good()); + REQUIRE(graph.create_execution_plans({fe::HeurMode_t::A}).is_good()); + REQUIRE(graph.check_support().is_good()); + REQUIRE(graph.build_plans().is_good()); + REQUIRE(graph.get_workspace_size() == 0); + + std::unordered_map, void *> variant_pack = { + {Q, nullptr}, {K, nullptr}, {V, nullptr}, {O, nullptr}}; + REQUIRE(graph.execute(nullptr, variant_pack, nullptr).is_good()); +} + +TEST_CASE("Mixed zero and non-zero element graph is rejected", "[graph][validate][zero_element]") { + namespace fe = cudnn_frontend; + fe::graph::Graph graph; + + graph.set_io_data_type(fe::DataType_t::HALF) + .set_intermediate_data_type(fe::DataType_t::FLOAT) + .set_compute_data_type(fe::DataType_t::FLOAT); + + // Matmul with a contracted dimension of 0: inputs are zero-element, but the + // output is not. This would require zero-filling the output, which cuDNN + // does not support; expect a clear validation error instead of a backend + // BAD_PARAM at build time. + auto A = graph.tensor(fe::graph::Tensor_attributes().set_name("A").set_dim({1, 4, 0}).set_stride({1, 1, 1})); + auto B = graph.tensor(fe::graph::Tensor_attributes().set_name("B").set_dim({1, 0, 8}).set_stride({1, 1, 1})); + + auto C = graph.matmul(A, B, fe::graph::Matmul_attributes()); + C->set_output(true); + + auto status = graph.validate(); + REQUIRE(status.get_code() == fe::error_code_t::GRAPH_NOT_SUPPORTED); + REQUIRE(status.get_message().find("zero-element") != std::string::npos); + REQUIRE_FALSE(graph.is_zero_element_graph()); } \ No newline at end of file diff --git a/test/python/test_zero_element_graph.py b/test/python/test_zero_element_graph.py new file mode 100644 index 000000000..55b5b8db6 --- /dev/null +++ b/test/python/test_zero_element_graph.py @@ -0,0 +1,150 @@ +""" +Tests for zero-element (no-op) graph support, e.g. SDPA with batch size 0. + +See https://github.com/NVIDIA/cudnn-frontend/issues/101: graphs whose output +tensors are all zero-element (a dimension of size 0) are treated as no-ops. +They validate and build successfully, report a workspace size of 0, and +execute() launches no work. +""" + +import cudnn +import pytest +import torch + +from test_utils import torch_fork_set_rng + + +def convert_to_cudnn_type(torch_type): + if torch_type == torch.float16: + return cudnn.data_type.HALF + elif torch_type == torch.bfloat16: + return cudnn.data_type.BFLOAT16 + elif torch_type == torch.float32: + return cudnn.data_type.FLOAT + else: + raise ValueError("Unsupported tensor data type.") + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_sdpa_batch_size_zero(cudnn_handle): + b, h, s_q, s_kv, d = 0, 10, 128, 16, 64 + + dtype = torch.float16 + q_gpu = torch.zeros(b, h, s_q, d, device="cuda", dtype=dtype) + k_gpu = torch.zeros(b, h, s_kv, d, device="cuda", dtype=dtype) + v_gpu = torch.zeros(b, h, s_kv, d, device="cuda", dtype=dtype) + o_gpu = torch.zeros(b, h, s_q, d, device="cuda", dtype=dtype) + + stream = torch.cuda.current_stream().cuda_stream + cudnn.set_stream(handle=cudnn_handle, stream=stream) + + graph = cudnn.pygraph( + handle=cudnn_handle, + io_data_type=convert_to_cudnn_type(dtype), + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + + q = graph.tensor_like(q_gpu) + k = graph.tensor_like(k_gpu) + v = graph.tensor_like(v_gpu) + + o, _ = graph.sdpa( + name="sdpa", + q=q, + k=k, + v=v, + generate_stats=False, + attn_scale=1.0 / (d**0.5), + ) + o.set_output(True).set_dim(o_gpu.size()).set_stride(o_gpu.stride()) + + graph.validate() + assert graph.is_zero_element_graph() + + graph.build_operation_graph() + graph.create_execution_plans([cudnn.heur_mode.A]) + graph.check_support() + graph.build_plans() + + assert graph.get_workspace_size() == 0 + workspace = torch.empty(graph.get_workspace_size(), device="cuda", dtype=torch.uint8) + + variant_pack = { + q: q_gpu, + k: k_gpu, + v: v_gpu, + o: o_gpu, + } + graph.execute(variant_pack, workspace, handle=cudnn_handle) + torch.cuda.synchronize() + + assert o_gpu.shape == (b, h, s_q, d) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_pointwise_zero_element(cudnn_handle): + dims = (0, 8, 16) + + dtype = torch.float32 + a_gpu = torch.zeros(*dims, device="cuda", dtype=dtype) + b_gpu = torch.zeros(*dims, device="cuda", dtype=dtype) + c_gpu = torch.zeros(*dims, device="cuda", dtype=dtype) + + stream = torch.cuda.current_stream().cuda_stream + cudnn.set_stream(handle=cudnn_handle, stream=stream) + + graph = cudnn.pygraph( + handle=cudnn_handle, + io_data_type=convert_to_cudnn_type(dtype), + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + + a = graph.tensor_like(a_gpu) + b = graph.tensor_like(b_gpu) + + c = graph.add(a=a, b=b) + c.set_output(True).set_dim(c_gpu.size()).set_stride(c_gpu.stride()) + + graph.build([cudnn.heur_mode.A]) + assert graph.is_zero_element_graph() + assert graph.get_workspace_size() == 0 + workspace = torch.empty(graph.get_workspace_size(), device="cuda", dtype=torch.uint8) + + variant_pack = {a: a_gpu, b: b_gpu, c: c_gpu} + graph.execute(variant_pack, workspace, handle=cudnn_handle) + torch.cuda.synchronize() + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_mixed_zero_element_graph_rejected(cudnn_handle): + # Matmul with a contracted dimension of 0: inputs are zero-element but the + # output is not. This would require zero-filling the output, which cuDNN + # does not support; expect a clear validation error. + stream = torch.cuda.current_stream().cuda_stream + cudnn.set_stream(handle=cudnn_handle, stream=stream) + + graph = cudnn.pygraph( + handle=cudnn_handle, + io_data_type=cudnn.data_type.HALF, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + + a = graph.tensor(name="A", dim=[1, 4, 0], stride=[1, 1, 1]) + b = graph.tensor(name="B", dim=[1, 0, 8], stride=[1, 1, 1]) + + c = graph.matmul(name="matmul", A=a, B=b) + c.set_output(True) + + with pytest.raises(cudnn.cudnnGraphNotSupportedError, match="zero-element"): + graph.validate() + assert not graph.is_zero_element_graph() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])