// Copyright © Advanced Micro Devices, Inc., or its affiliates. // SPDX-License-Identifier: MIT #include #include #include #include #include #include #include #include #include #include #include #include #include "test_plugins/TestPluginConstants.hpp" using namespace hipdnn_frontend; using namespace hipdnn_frontend::graph; using namespace hipdnn_tests::constants; using hipdnn_tests::toVec; using DataTypeSdk = hipdnn_flatbuffers_sdk::data_objects::DataType; using NodeAttrType = hipdnn_flatbuffers_sdk::data_objects::NodeAttributes; using DiagonalAlignmentSdk = hipdnn_flatbuffers_sdk::data_objects::DiagonalAlignment; namespace { // Exposes protected Graph methods for testing class TestableGraph : public Graph { public: using Graph::build_operation_graph_via_descriptors; using Graph::get_raw_graph_descriptor; }; // Lowers a frontend graph via build_operation_graph_via_descriptors, then // retrieves the serialized graph and deserializes it for verification. class IntegrationSdpaBwdDescriptorLowering : public ::testing::Test { protected: void SetUp() override { SKIP_IF_NO_DEVICES(); ASSERT_EQ(hipInit(0), hipSuccess); const std::array paths = {hipdnn_tests::plugin_constants::testGoodPluginPath().c_str()}; ASSERT_EQ(hipdnnSetEnginePluginPaths_ext( paths.size(), paths.data(), HIPDNN_PLUGIN_LOADING_ABSOLUTE), HIPDNN_STATUS_SUCCESS); ASSERT_EQ(hipdnnCreate(&_handle), HIPDNN_STATUS_SUCCESS); } void TearDown() override { if(_handle != nullptr) { hipdnnDestroy(_handle); } } hipdnnHandle_t _handle = nullptr; }; // Builds an SDPA backward graph via the frontend API, lowers it to the backend // via build_operation_graph_via_descriptors, retrieves the serialized graph, // and verifies ALL tensor and operation attributes match the values set // in the frontend. TEST_F(IntegrationSdpaBwdDescriptorLowering, SdpaBwdGraphRoundTrip) { auto graph = std::make_shared(); graph->set_name("TestSdpaBwdGraph") .set_io_data_type(DataType::FLOAT) .set_intermediate_data_type(DataType::FLOAT) .set_compute_data_type(DataType::FLOAT); auto q = std::make_shared(); q->set_uid(K_SDPA_BWD_TENSOR_Q_UID).set_name("Q").set_data_type(DataType::FLOAT); q->set_dim(toVec(K_SDPA_BWD_TENSOR_Q_DIMS)).set_stride(toVec(K_SDPA_BWD_TENSOR_Q_STRIDES)); auto k = std::make_shared(); k->set_uid(K_SDPA_BWD_TENSOR_K_UID).set_name("K").set_data_type(DataType::FLOAT); k->set_dim(toVec(K_SDPA_BWD_TENSOR_K_DIMS)).set_stride(toVec(K_SDPA_BWD_TENSOR_K_STRIDES)); auto v = std::make_shared(); v->set_uid(K_SDPA_BWD_TENSOR_V_UID).set_name("V").set_data_type(DataType::FLOAT); v->set_dim(toVec(K_SDPA_BWD_TENSOR_V_DIMS)).set_stride(toVec(K_SDPA_BWD_TENSOR_V_STRIDES)); auto o = std::make_shared(); o->set_uid(K_SDPA_BWD_TENSOR_O_UID).set_name("O").set_data_type(DataType::FLOAT); o->set_dim(toVec(K_SDPA_BWD_TENSOR_O_DIMS)).set_stride(toVec(K_SDPA_BWD_TENSOR_O_STRIDES)); auto dO = std::make_shared(); dO->set_uid(K_SDPA_BWD_TENSOR_DO_UID).set_name("dO").set_data_type(DataType::FLOAT); dO->set_dim(toVec(K_SDPA_BWD_TENSOR_DO_DIMS)).set_stride(toVec(K_SDPA_BWD_TENSOR_DO_STRIDES)); auto stats = std::make_shared(); stats->set_uid(K_SDPA_BWD_TENSOR_STATS_UID).set_name("Stats").set_data_type(DataType::FLOAT); stats->set_dim(toVec(K_SDPA_BWD_TENSOR_STATS_DIMS)) .set_stride(toVec(K_SDPA_BWD_TENSOR_STATS_STRIDES)); SdpaBackwardAttributes sdpaAttrs; sdpaAttrs.set_name("sdpa_bwd_op"); auto [dq, dk, dv] = graph->sdpa_backward(q, k, v, o, dO, stats, std::move(sdpaAttrs)); dq->set_uid(K_SDPA_BWD_TENSOR_DQ_UID).set_output(true).set_name("dQ"); dk->set_uid(K_SDPA_BWD_TENSOR_DK_UID).set_output(true).set_name("dK"); dv->set_uid(K_SDPA_BWD_TENSOR_DV_UID).set_output(true).set_name("dV"); // -- Validate and lower -- auto result = graph->validate(); ASSERT_EQ(result.code, ErrorCode::OK) << result.err_msg; result = graph->build_operation_graph_via_descriptors(_handle); ASSERT_EQ(result.code, ErrorCode::OK) << result.err_msg; // -- Retrieve serialized graph -- auto rawDesc = graph->get_raw_graph_descriptor(); ASSERT_NE(rawDesc, nullptr); size_t serializedSize = 0; ASSERT_EQ(hipdnnBackendGetSerializedBinaryGraph_ext(rawDesc, 0, &serializedSize, nullptr), HIPDNN_STATUS_SUCCESS); ASSERT_GT(serializedSize, 0u); std::vector serializedData(serializedSize); ASSERT_EQ(hipdnnBackendGetSerializedBinaryGraph_ext( rawDesc, serializedSize, &serializedSize, serializedData.data()), HIPDNN_STATUS_SUCCESS); // -- Deserialize into GraphT -- auto graphFb = hipdnn_flatbuffers_sdk::data_objects::GetGraph(serializedData.data()); ASSERT_NE(graphFb, nullptr); hipdnn_flatbuffers_sdk::data_objects::GraphT graphT; graphFb->UnPackTo(&graphT); // -- Verify graph-level attributes -- EXPECT_EQ(graphT.compute_data_type, DataTypeSdk::FLOAT); EXPECT_EQ(graphT.intermediate_data_type, DataTypeSdk::FLOAT); EXPECT_EQ(graphT.io_data_type, DataTypeSdk::FLOAT); // -- Verify tensors (Q, K, V, O, dO, Stats, dQ, dK, dV = 9 tensors) -- ASSERT_EQ(graphT.tensors.size(), 9u); std::unordered_map tensorMap; for(const auto& t : graphT.tensors) { tensorMap[t->uid] = t.get(); } // Verify Q tensor - UID, name, data type, dims, strides ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_Q_UID), 0u); auto* qT = tensorMap[K_SDPA_BWD_TENSOR_Q_UID]; EXPECT_EQ(qT->name, "Q"); EXPECT_EQ(qT->data_type, DataTypeSdk::FLOAT); EXPECT_EQ(qT->dims, toVec(K_SDPA_BWD_TENSOR_Q_DIMS)); EXPECT_EQ(qT->strides, toVec(K_SDPA_BWD_TENSOR_Q_STRIDES)); EXPECT_FALSE(qT->virtual_); // Verify K tensor - UID, name, data type, dims, strides ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_K_UID), 0u); auto* kT = tensorMap[K_SDPA_BWD_TENSOR_K_UID]; EXPECT_EQ(kT->name, "K"); EXPECT_EQ(kT->data_type, DataTypeSdk::FLOAT); EXPECT_EQ(kT->dims, toVec(K_SDPA_BWD_TENSOR_K_DIMS)); EXPECT_EQ(kT->strides, toVec(K_SDPA_BWD_TENSOR_K_STRIDES)); EXPECT_FALSE(kT->virtual_); // Verify V tensor - UID, name, data type, dims, strides ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_V_UID), 0u); auto* vT = tensorMap[K_SDPA_BWD_TENSOR_V_UID]; EXPECT_EQ(vT->name, "V"); EXPECT_EQ(vT->data_type, DataTypeSdk::FLOAT); EXPECT_EQ(vT->dims, toVec(K_SDPA_BWD_TENSOR_V_DIMS)); EXPECT_EQ(vT->strides, toVec(K_SDPA_BWD_TENSOR_V_STRIDES)); EXPECT_FALSE(vT->virtual_); // Verify O tensor - UID, name, data type, dims, strides ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_O_UID), 0u); auto* oT = tensorMap[K_SDPA_BWD_TENSOR_O_UID]; EXPECT_EQ(oT->name, "O"); EXPECT_EQ(oT->data_type, DataTypeSdk::FLOAT); EXPECT_EQ(oT->dims, toVec(K_SDPA_BWD_TENSOR_O_DIMS)); EXPECT_EQ(oT->strides, toVec(K_SDPA_BWD_TENSOR_O_STRIDES)); EXPECT_FALSE(oT->virtual_); // Verify dO tensor - UID, name, data type, dims, strides ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_DO_UID), 0u); auto* doT = tensorMap[K_SDPA_BWD_TENSOR_DO_UID]; EXPECT_EQ(doT->name, "dO"); EXPECT_EQ(doT->data_type, DataTypeSdk::FLOAT); EXPECT_EQ(doT->dims, toVec(K_SDPA_BWD_TENSOR_DO_DIMS)); EXPECT_EQ(doT->strides, toVec(K_SDPA_BWD_TENSOR_DO_STRIDES)); EXPECT_FALSE(doT->virtual_); // Verify Stats tensor - UID, name, data type, dims, strides ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_STATS_UID), 0u); auto* statsT = tensorMap[K_SDPA_BWD_TENSOR_STATS_UID]; EXPECT_EQ(statsT->name, "Stats"); EXPECT_EQ(statsT->data_type, DataTypeSdk::FLOAT); EXPECT_EQ(statsT->dims, toVec(K_SDPA_BWD_TENSOR_STATS_DIMS)); EXPECT_EQ(statsT->strides, toVec(K_SDPA_BWD_TENSOR_STATS_STRIDES)); EXPECT_FALSE(statsT->virtual_); // Verify dQ tensor (output) - UID, name, data type, dims, strides ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_DQ_UID), 0u); auto* dqT = tensorMap[K_SDPA_BWD_TENSOR_DQ_UID]; EXPECT_EQ(dqT->name, "dQ"); EXPECT_EQ(dqT->data_type, DataTypeSdk::FLOAT); EXPECT_EQ(dqT->dims, toVec(K_SDPA_BWD_TENSOR_DQ_DIMS)); EXPECT_EQ(dqT->strides, toVec(K_SDPA_BWD_TENSOR_DQ_STRIDES)); EXPECT_FALSE(dqT->virtual_); // Verify dK tensor (output) - UID, name, data type, dims, strides ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_DK_UID), 0u); auto* dkT = tensorMap[K_SDPA_BWD_TENSOR_DK_UID]; EXPECT_EQ(dkT->name, "dK"); EXPECT_EQ(dkT->data_type, DataTypeSdk::FLOAT); EXPECT_EQ(dkT->dims, toVec(K_SDPA_BWD_TENSOR_DK_DIMS)); EXPECT_EQ(dkT->strides, toVec(K_SDPA_BWD_TENSOR_DK_STRIDES)); EXPECT_FALSE(dkT->virtual_); // Verify dV tensor (output) - UID, name, data type, dims, strides ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_DV_UID), 0u); auto* dvT = tensorMap[K_SDPA_BWD_TENSOR_DV_UID]; EXPECT_EQ(dvT->name, "dV"); EXPECT_EQ(dvT->data_type, DataTypeSdk::FLOAT); EXPECT_EQ(dvT->dims, toVec(K_SDPA_BWD_TENSOR_DV_DIMS)); EXPECT_EQ(dvT->strides, toVec(K_SDPA_BWD_TENSOR_DV_STRIDES)); EXPECT_FALSE(dvT->virtual_); // -- Verify SDPA bprop operation node -- ASSERT_EQ(graphT.nodes.size(), 1u); auto& node = graphT.nodes[0]; EXPECT_EQ(node->name, "sdpa_bwd_op"); EXPECT_EQ(node->compute_data_type, DataTypeSdk::FLOAT); EXPECT_EQ(node->attributes.type, NodeAttrType::SdpaBackwardAttributes); auto* sdpa = node->attributes.AsSdpaBackwardAttributes(); ASSERT_NE(sdpa, nullptr); // Verify required tensor UID references in operation node EXPECT_EQ(sdpa->q_tensor_uid, K_SDPA_BWD_TENSOR_Q_UID); EXPECT_EQ(sdpa->k_tensor_uid, K_SDPA_BWD_TENSOR_K_UID); EXPECT_EQ(sdpa->v_tensor_uid, K_SDPA_BWD_TENSOR_V_UID); EXPECT_EQ(sdpa->o_tensor_uid, K_SDPA_BWD_TENSOR_O_UID); EXPECT_EQ(sdpa->do_tensor_uid, K_SDPA_BWD_TENSOR_DO_UID); EXPECT_EQ(sdpa->stats_tensor_uid, K_SDPA_BWD_TENSOR_STATS_UID); EXPECT_EQ(sdpa->dq_tensor_uid, K_SDPA_BWD_TENSOR_DQ_UID); EXPECT_EQ(sdpa->dk_tensor_uid, K_SDPA_BWD_TENSOR_DK_UID); EXPECT_EQ(sdpa->dv_tensor_uid, K_SDPA_BWD_TENSOR_DV_UID); // Verify optional tensor UIDs are absent when not set EXPECT_FALSE(sdpa->scale_tensor_uid.has_value()); EXPECT_FALSE(sdpa->attn_mask_tensor_uid.has_value()); EXPECT_FALSE(sdpa->seq_len_q_tensor_uid.has_value()); EXPECT_FALSE(sdpa->seq_len_kv_tensor_uid.has_value()); EXPECT_FALSE(sdpa->seed_tensor_uid.has_value()); EXPECT_FALSE(sdpa->offset_tensor_uid.has_value()); EXPECT_FALSE(sdpa->dropout_mask_tensor_uid.has_value()); EXPECT_FALSE(sdpa->dropout_scale_tensor_uid.has_value()); EXPECT_FALSE(sdpa->dropout_scale_inv_tensor_uid.has_value()); EXPECT_FALSE(sdpa->dbias_tensor_uid.has_value()); // Verify all default scalar/boolean/enum values EXPECT_EQ(sdpa->diagonal_alignment, DiagonalAlignmentSdk::TOP_LEFT); EXPECT_FALSE(sdpa->alibi_mask); EXPECT_FALSE(sdpa->padding_mask); EXPECT_FALSE(sdpa->causal_mask); EXPECT_FALSE(sdpa->causal_mask_bottom_right); EXPECT_FALSE(sdpa->dropout_probability.has_value()); EXPECT_FALSE(sdpa->attn_scale_value.has_value()); EXPECT_FALSE(sdpa->left_bound.has_value()); EXPECT_FALSE(sdpa->right_bound.has_value()); } // Sets ALL optional tensors and all non-default scalar/boolean/enum values, // then verifies every tensor property and every field survives the round-trip. TEST_F(IntegrationSdpaBwdDescriptorLowering, SdpaBwdWithAllOptionalTensorsAndScalars) { auto graph = std::make_shared(); graph->set_name("TestSdpaBwdAllOptionals") .set_io_data_type(DataType::FLOAT) .set_intermediate_data_type(DataType::FLOAT) .set_compute_data_type(DataType::FLOAT); // -- Required input tensors -- auto q = std::make_shared(); q->set_uid(K_SDPA_BWD_TENSOR_Q_UID).set_name("Q").set_data_type(DataType::FLOAT); q->set_dim(toVec(K_SDPA_BWD_TENSOR_Q_DIMS)).set_stride(toVec(K_SDPA_BWD_TENSOR_Q_STRIDES)); auto k = std::make_shared(); k->set_uid(K_SDPA_BWD_TENSOR_K_UID).set_name("K").set_data_type(DataType::FLOAT); k->set_dim(toVec(K_SDPA_BWD_TENSOR_K_DIMS)).set_stride(toVec(K_SDPA_BWD_TENSOR_K_STRIDES)); auto v = std::make_shared(); v->set_uid(K_SDPA_BWD_TENSOR_V_UID).set_name("V").set_data_type(DataType::FLOAT); v->set_dim(toVec(K_SDPA_BWD_TENSOR_V_DIMS)).set_stride(toVec(K_SDPA_BWD_TENSOR_V_STRIDES)); auto o = std::make_shared(); o->set_uid(K_SDPA_BWD_TENSOR_O_UID).set_name("O").set_data_type(DataType::FLOAT); o->set_dim(toVec(K_SDPA_BWD_TENSOR_O_DIMS)).set_stride(toVec(K_SDPA_BWD_TENSOR_O_STRIDES)); auto dO = std::make_shared(); dO->set_uid(K_SDPA_BWD_TENSOR_DO_UID).set_name("dO").set_data_type(DataType::FLOAT); dO->set_dim(toVec(K_SDPA_BWD_TENSOR_DO_DIMS)).set_stride(toVec(K_SDPA_BWD_TENSOR_DO_STRIDES)); auto statsTensor = std::make_shared(); statsTensor->set_uid(K_SDPA_BWD_TENSOR_STATS_UID) .set_name("Stats") .set_data_type(DataType::FLOAT); statsTensor->set_dim(toVec(K_SDPA_BWD_TENSOR_STATS_DIMS)) .set_stride(toVec(K_SDPA_BWD_TENSOR_STATS_STRIDES)); // -- ALL optional input tensors -- auto scale = std::make_shared(); scale->set_uid(K_SDPA_BWD_TENSOR_SCALE_UID).set_name("SCALE"); scale->set_value(0.125f); auto attnMask = std::make_shared(); attnMask->set_uid(K_SDPA_BWD_TENSOR_ATTN_MASK_UID) .set_name("ATTN_MASK") .set_data_type(DataType::FLOAT); attnMask->set_dim(toVec(K_SDPA_BWD_TENSOR_ATTN_MASK_DIMS)); attnMask->set_stride(toVec(K_SDPA_BWD_TENSOR_ATTN_MASK_STRIDES)); auto seqLenQ = std::make_shared(); seqLenQ->set_uid(K_SDPA_BWD_TENSOR_SEQ_LEN_Q_UID) .set_name("SEQ_LEN_Q") .set_data_type(DataType::INT32); seqLenQ->set_dim(toVec(K_SDPA_BWD_TENSOR_SEQ_LEN_DIMS)); seqLenQ->set_stride(toVec(K_SDPA_BWD_TENSOR_SEQ_LEN_STRIDES)); auto seqLenKv = std::make_shared(); seqLenKv->set_uid(K_SDPA_BWD_TENSOR_SEQ_LEN_KV_UID) .set_name("SEQ_LEN_KV") .set_data_type(DataType::INT32); seqLenKv->set_dim(toVec(K_SDPA_BWD_TENSOR_SEQ_LEN_DIMS)); seqLenKv->set_stride(toVec(K_SDPA_BWD_TENSOR_SEQ_LEN_STRIDES)); auto seed = std::make_shared(); seed->set_uid(K_SDPA_BWD_TENSOR_SEED_UID).set_name("SEED"); seed->set_value(int32_t{42}); auto offset = std::make_shared(); offset->set_uid(K_SDPA_BWD_TENSOR_OFFSET_UID).set_name("OFFSET"); offset->set_value(int32_t{0}); auto dropoutMask = std::make_shared(); dropoutMask->set_uid(K_SDPA_BWD_TENSOR_DROPOUT_MASK_UID) .set_name("DROPOUT_MASK") .set_data_type(DataType::FLOAT); dropoutMask->set_dim(toVec(K_SDPA_BWD_TENSOR_DROPOUT_MASK_DIMS)); dropoutMask->set_stride(toVec(K_SDPA_BWD_TENSOR_DROPOUT_MASK_STRIDES)); auto dropoutScale = std::make_shared(); dropoutScale->set_uid(K_SDPA_BWD_TENSOR_DROPOUT_SCALE_UID).set_name("DROPOUT_SCALE"); dropoutScale->set_value(1.0f / (1.0f - 0.1f)); auto dropoutScaleInv = std::make_shared(); dropoutScaleInv->set_uid(K_SDPA_BWD_TENSOR_DROPOUT_SCALE_INV_UID).set_name("DROPOUT_SCALE_INV"); dropoutScaleInv->set_value(1.0f - 0.1f); // -- Optional output tensor: dBias (must be set on attrs before graph call) -- auto dbias = std::make_shared(); dbias->set_uid(K_SDPA_BWD_TENSOR_DBIAS_UID) .set_name("DBIAS") .set_data_type(DataType::FLOAT) .set_output(true); dbias->set_dim(toVec(K_SDPA_BWD_TENSOR_DBIAS_DIMS)); dbias->set_stride(toVec(K_SDPA_BWD_TENSOR_DBIAS_STRIDES)); SdpaBackwardAttributes sdpaAttrs; sdpaAttrs.set_name("sdpa_bwd_all_optionals"); sdpaAttrs.set_attn_scale(scale); sdpaAttrs.set_bias(attnMask); sdpaAttrs.set_seq_len_q(seqLenQ); sdpaAttrs.set_seq_len_kv(seqLenKv); sdpaAttrs.set_dropout(0.1f, seed, offset); sdpaAttrs.set_dropout_mask(dropoutMask); sdpaAttrs.set_dropout_scale(dropoutScale); sdpaAttrs.set_dropout_scale_inv(dropoutScaleInv); sdpaAttrs.set_dbias(dbias); sdpaAttrs.set_alibi_mask(true); sdpaAttrs.set_padding_mask(true); sdpaAttrs.set_causal_mask(true); sdpaAttrs.set_causal_mask_bottom_right(true); sdpaAttrs.set_attn_scale(0.125f); sdpaAttrs.set_diagonal_band_left_bound(0); sdpaAttrs.set_diagonal_band_right_bound(128); sdpaAttrs.set_diagonal_alignment(DiagonalAlignment::BOTTOM_RIGHT); auto [dq, dk, dv] = graph->sdpa_backward(q, k, v, o, dO, statsTensor, std::move(sdpaAttrs)); dq->set_uid(K_SDPA_BWD_TENSOR_DQ_UID).set_output(true).set_name("dQ"); dk->set_uid(K_SDPA_BWD_TENSOR_DK_UID).set_output(true).set_name("dK"); dv->set_uid(K_SDPA_BWD_TENSOR_DV_UID).set_output(true).set_name("dV"); // -- Validate and lower -- auto result = graph->validate(); ASSERT_EQ(result.code, ErrorCode::OK) << result.err_msg; result = graph->build_operation_graph_via_descriptors(_handle); ASSERT_EQ(result.code, ErrorCode::OK) << result.err_msg; // -- Retrieve serialized graph -- auto rawDesc = graph->get_raw_graph_descriptor(); ASSERT_NE(rawDesc, nullptr); size_t serializedSize = 0; ASSERT_EQ(hipdnnBackendGetSerializedBinaryGraph_ext(rawDesc, 0, &serializedSize, nullptr), HIPDNN_STATUS_SUCCESS); ASSERT_GT(serializedSize, 0u); std::vector serializedData(serializedSize); ASSERT_EQ(hipdnnBackendGetSerializedBinaryGraph_ext( rawDesc, serializedSize, &serializedSize, serializedData.data()), HIPDNN_STATUS_SUCCESS); // -- Deserialize into GraphT -- hipdnn_flatbuffers_sdk::data_objects::GraphT graphT; hipdnn_flatbuffers_sdk::data_objects::GetGraph(serializedData.data())->UnPackTo(&graphT); // 9 required + 10 optional input/output tensors (including scale) = 19 ASSERT_EQ(graphT.tensors.size(), 19u); std::unordered_map tensorMap; for(const auto& t : graphT.tensors) { tensorMap[t->uid] = t.get(); } // -- Verify all tensor UIDs are present -- // Required ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_Q_UID), 0u); ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_K_UID), 0u); ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_V_UID), 0u); ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_O_UID), 0u); ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_DO_UID), 0u); ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_STATS_UID), 0u); ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_DQ_UID), 0u); ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_DK_UID), 0u); ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_DV_UID), 0u); // Optional tensors (all present when all optionals are set) ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_SCALE_UID), 0u); ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_ATTN_MASK_UID), 0u); ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_SEQ_LEN_Q_UID), 0u); ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_SEQ_LEN_KV_UID), 0u); ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_SEED_UID), 0u); ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_OFFSET_UID), 0u); ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_DROPOUT_MASK_UID), 0u); ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_DROPOUT_SCALE_UID), 0u); ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_DROPOUT_SCALE_INV_UID), 0u); ASSERT_NE(tensorMap.count(K_SDPA_BWD_TENSOR_DBIAS_UID), 0u); // -- Verify optional tensor properties -- auto* scaleT = tensorMap[K_SDPA_BWD_TENSOR_SCALE_UID]; EXPECT_EQ(scaleT->name, "SCALE"); EXPECT_EQ(scaleT->dims, toVec(K_SDPA_BWD_TENSOR_SCALAR_DIMS)); EXPECT_EQ(scaleT->strides, toVec(K_SDPA_BWD_TENSOR_SCALAR_STRIDES)); auto* attnMaskT = tensorMap[K_SDPA_BWD_TENSOR_ATTN_MASK_UID]; EXPECT_EQ(attnMaskT->name, "ATTN_MASK"); EXPECT_EQ(attnMaskT->data_type, DataTypeSdk::FLOAT); EXPECT_EQ(attnMaskT->dims, toVec(K_SDPA_BWD_TENSOR_ATTN_MASK_DIMS)); EXPECT_EQ(attnMaskT->strides, toVec(K_SDPA_BWD_TENSOR_ATTN_MASK_STRIDES)); auto* seqLenQT = tensorMap[K_SDPA_BWD_TENSOR_SEQ_LEN_Q_UID]; EXPECT_EQ(seqLenQT->name, "SEQ_LEN_Q"); EXPECT_EQ(seqLenQT->data_type, DataTypeSdk::INT32); EXPECT_EQ(seqLenQT->dims, toVec(K_SDPA_BWD_TENSOR_SEQ_LEN_DIMS)); EXPECT_EQ(seqLenQT->strides, toVec(K_SDPA_BWD_TENSOR_SEQ_LEN_STRIDES)); auto* seqLenKvT = tensorMap[K_SDPA_BWD_TENSOR_SEQ_LEN_KV_UID]; EXPECT_EQ(seqLenKvT->name, "SEQ_LEN_KV"); EXPECT_EQ(seqLenKvT->data_type, DataTypeSdk::INT32); EXPECT_EQ(seqLenKvT->dims, toVec(K_SDPA_BWD_TENSOR_SEQ_LEN_DIMS)); EXPECT_EQ(seqLenKvT->strides, toVec(K_SDPA_BWD_TENSOR_SEQ_LEN_STRIDES)); auto* seedT = tensorMap[K_SDPA_BWD_TENSOR_SEED_UID]; EXPECT_EQ(seedT->name, "SEED"); EXPECT_EQ(seedT->data_type, DataTypeSdk::INT32); EXPECT_EQ(seedT->dims, toVec(K_SDPA_BWD_TENSOR_SCALAR_DIMS)); EXPECT_EQ(seedT->strides, toVec(K_SDPA_BWD_TENSOR_SCALAR_STRIDES)); auto* offsetT = tensorMap[K_SDPA_BWD_TENSOR_OFFSET_UID]; EXPECT_EQ(offsetT->name, "OFFSET"); EXPECT_EQ(offsetT->data_type, DataTypeSdk::INT32); EXPECT_EQ(offsetT->dims, toVec(K_SDPA_BWD_TENSOR_SCALAR_DIMS)); EXPECT_EQ(offsetT->strides, toVec(K_SDPA_BWD_TENSOR_SCALAR_STRIDES)); auto* dropoutMaskT = tensorMap[K_SDPA_BWD_TENSOR_DROPOUT_MASK_UID]; EXPECT_EQ(dropoutMaskT->name, "DROPOUT_MASK"); EXPECT_EQ(dropoutMaskT->data_type, DataTypeSdk::FLOAT); EXPECT_EQ(dropoutMaskT->dims, toVec(K_SDPA_BWD_TENSOR_DROPOUT_MASK_DIMS)); EXPECT_EQ(dropoutMaskT->strides, toVec(K_SDPA_BWD_TENSOR_DROPOUT_MASK_STRIDES)); auto* dropoutScaleT = tensorMap[K_SDPA_BWD_TENSOR_DROPOUT_SCALE_UID]; EXPECT_EQ(dropoutScaleT->name, "DROPOUT_SCALE"); EXPECT_EQ(dropoutScaleT->data_type, DataTypeSdk::FLOAT); EXPECT_EQ(dropoutScaleT->dims, toVec(K_SDPA_BWD_TENSOR_SCALAR_DIMS)); EXPECT_EQ(dropoutScaleT->strides, toVec(K_SDPA_BWD_TENSOR_SCALAR_STRIDES)); auto* dropoutScaleInvT = tensorMap[K_SDPA_BWD_TENSOR_DROPOUT_SCALE_INV_UID]; EXPECT_EQ(dropoutScaleInvT->name, "DROPOUT_SCALE_INV"); EXPECT_EQ(dropoutScaleInvT->data_type, DataTypeSdk::FLOAT); EXPECT_EQ(dropoutScaleInvT->dims, toVec(K_SDPA_BWD_TENSOR_SCALAR_DIMS)); EXPECT_EQ(dropoutScaleInvT->strides, toVec(K_SDPA_BWD_TENSOR_SCALAR_STRIDES)); auto* dbiasT = tensorMap[K_SDPA_BWD_TENSOR_DBIAS_UID]; EXPECT_EQ(dbiasT->name, "DBIAS"); EXPECT_EQ(dbiasT->data_type, DataTypeSdk::FLOAT); EXPECT_EQ(dbiasT->dims, toVec(K_SDPA_BWD_TENSOR_DBIAS_DIMS)); EXPECT_EQ(dbiasT->strides, toVec(K_SDPA_BWD_TENSOR_DBIAS_STRIDES)); // -- Verify SDPA bprop node attributes -- ASSERT_EQ(graphT.nodes.size(), 1u); auto& node = graphT.nodes[0]; EXPECT_EQ(node->name, "sdpa_bwd_all_optionals"); EXPECT_EQ(node->compute_data_type, DataTypeSdk::FLOAT); EXPECT_EQ(node->attributes.type, NodeAttrType::SdpaBackwardAttributes); auto* sdpa = node->attributes.AsSdpaBackwardAttributes(); ASSERT_NE(sdpa, nullptr); // Required tensor UIDs EXPECT_EQ(sdpa->q_tensor_uid, K_SDPA_BWD_TENSOR_Q_UID); EXPECT_EQ(sdpa->k_tensor_uid, K_SDPA_BWD_TENSOR_K_UID); EXPECT_EQ(sdpa->v_tensor_uid, K_SDPA_BWD_TENSOR_V_UID); EXPECT_EQ(sdpa->o_tensor_uid, K_SDPA_BWD_TENSOR_O_UID); EXPECT_EQ(sdpa->do_tensor_uid, K_SDPA_BWD_TENSOR_DO_UID); EXPECT_EQ(sdpa->stats_tensor_uid, K_SDPA_BWD_TENSOR_STATS_UID); EXPECT_EQ(sdpa->dq_tensor_uid, K_SDPA_BWD_TENSOR_DQ_UID); EXPECT_EQ(sdpa->dk_tensor_uid, K_SDPA_BWD_TENSOR_DK_UID); EXPECT_EQ(sdpa->dv_tensor_uid, K_SDPA_BWD_TENSOR_DV_UID); // ALL optional tensor UIDs should be present ASSERT_TRUE(sdpa->scale_tensor_uid.has_value()); EXPECT_EQ(sdpa->scale_tensor_uid.value(), K_SDPA_BWD_TENSOR_SCALE_UID); ASSERT_TRUE(sdpa->attn_mask_tensor_uid.has_value()); EXPECT_EQ(sdpa->attn_mask_tensor_uid.value(), K_SDPA_BWD_TENSOR_ATTN_MASK_UID); ASSERT_TRUE(sdpa->seq_len_q_tensor_uid.has_value()); EXPECT_EQ(sdpa->seq_len_q_tensor_uid.value(), K_SDPA_BWD_TENSOR_SEQ_LEN_Q_UID); ASSERT_TRUE(sdpa->seq_len_kv_tensor_uid.has_value()); EXPECT_EQ(sdpa->seq_len_kv_tensor_uid.value(), K_SDPA_BWD_TENSOR_SEQ_LEN_KV_UID); ASSERT_TRUE(sdpa->seed_tensor_uid.has_value()); EXPECT_EQ(sdpa->seed_tensor_uid.value(), K_SDPA_BWD_TENSOR_SEED_UID); ASSERT_TRUE(sdpa->offset_tensor_uid.has_value()); EXPECT_EQ(sdpa->offset_tensor_uid.value(), K_SDPA_BWD_TENSOR_OFFSET_UID); ASSERT_TRUE(sdpa->dropout_mask_tensor_uid.has_value()); EXPECT_EQ(sdpa->dropout_mask_tensor_uid.value(), K_SDPA_BWD_TENSOR_DROPOUT_MASK_UID); ASSERT_TRUE(sdpa->dropout_scale_tensor_uid.has_value()); EXPECT_EQ(sdpa->dropout_scale_tensor_uid.value(), K_SDPA_BWD_TENSOR_DROPOUT_SCALE_UID); ASSERT_TRUE(sdpa->dropout_scale_inv_tensor_uid.has_value()); EXPECT_EQ(sdpa->dropout_scale_inv_tensor_uid.value(), K_SDPA_BWD_TENSOR_DROPOUT_SCALE_INV_UID); ASSERT_TRUE(sdpa->dbias_tensor_uid.has_value()); EXPECT_EQ(sdpa->dbias_tensor_uid.value(), K_SDPA_BWD_TENSOR_DBIAS_UID); // ALL boolean attributes set to true EXPECT_TRUE(sdpa->alibi_mask); EXPECT_TRUE(sdpa->padding_mask); EXPECT_TRUE(sdpa->causal_mask); EXPECT_TRUE(sdpa->causal_mask_bottom_right); // Float scalars ASSERT_TRUE(sdpa->dropout_probability.has_value()); EXPECT_FLOAT_EQ(sdpa->dropout_probability.value(), 0.1f); ASSERT_TRUE(sdpa->attn_scale_value.has_value()); EXPECT_FLOAT_EQ(sdpa->attn_scale_value.value(), 0.125f); // Integer scalars ASSERT_TRUE(sdpa->left_bound.has_value()); EXPECT_EQ(sdpa->left_bound.value(), 0); ASSERT_TRUE(sdpa->right_bound.has_value()); EXPECT_EQ(sdpa->right_bound.value(), 128); // Enum values EXPECT_EQ(sdpa->diagonal_alignment, DiagonalAlignmentSdk::BOTTOM_RIGHT); } // Verifies that tensor UIDs auto-assigned by the frontend are preserved // through the lowering round-trip. Tensors are created WITHOUT explicit UIDs. TEST_F(IntegrationSdpaBwdDescriptorLowering, AutoAssignedUidsPreservedInRoundTrip) { auto graph = std::make_shared(); graph->set_name("AutoUidSdpaBwdGraph") .set_io_data_type(DataType::FLOAT) .set_intermediate_data_type(DataType::FLOAT) .set_compute_data_type(DataType::FLOAT); // Create tensors WITHOUT explicit UIDs auto q = std::make_shared(); q->set_name("Q").set_data_type(DataType::FLOAT); q->set_dim(toVec(K_SDPA_BWD_TENSOR_Q_DIMS)).set_stride(toVec(K_SDPA_BWD_TENSOR_Q_STRIDES)); auto k = std::make_shared(); k->set_name("K").set_data_type(DataType::FLOAT); k->set_dim(toVec(K_SDPA_BWD_TENSOR_K_DIMS)).set_stride(toVec(K_SDPA_BWD_TENSOR_K_STRIDES)); auto v = std::make_shared(); v->set_name("V").set_data_type(DataType::FLOAT); v->set_dim(toVec(K_SDPA_BWD_TENSOR_V_DIMS)).set_stride(toVec(K_SDPA_BWD_TENSOR_V_STRIDES)); auto o = std::make_shared(); o->set_name("O").set_data_type(DataType::FLOAT); o->set_dim(toVec(K_SDPA_BWD_TENSOR_O_DIMS)).set_stride(toVec(K_SDPA_BWD_TENSOR_O_STRIDES)); auto dO = std::make_shared(); dO->set_name("dO").set_data_type(DataType::FLOAT); dO->set_dim(toVec(K_SDPA_BWD_TENSOR_DO_DIMS)).set_stride(toVec(K_SDPA_BWD_TENSOR_DO_STRIDES)); auto stats = std::make_shared(); stats->set_name("Stats").set_data_type(DataType::FLOAT); stats->set_dim(toVec(K_SDPA_BWD_TENSOR_STATS_DIMS)) .set_stride(toVec(K_SDPA_BWD_TENSOR_STATS_STRIDES)); SdpaBackwardAttributes sdpaAttrs; auto [dq, dk, dv] = graph->sdpa_backward(q, k, v, o, dO, stats, std::move(sdpaAttrs)); dq->set_output(true); dk->set_output(true); dv->set_output(true); auto result = graph->validate(); ASSERT_EQ(result.code, ErrorCode::OK) << result.err_msg; result = graph->build_operation_graph_via_descriptors(_handle); ASSERT_EQ(result.code, ErrorCode::OK) << result.err_msg; // Retrieve serialized graph auto rawDesc = graph->get_raw_graph_descriptor(); size_t serializedSize = 0; ASSERT_EQ(hipdnnBackendGetSerializedBinaryGraph_ext(rawDesc, 0, &serializedSize, nullptr), HIPDNN_STATUS_SUCCESS); ASSERT_GT(serializedSize, 0u); std::vector serializedData(serializedSize); ASSERT_EQ(hipdnnBackendGetSerializedBinaryGraph_ext( rawDesc, serializedSize, &serializedSize, serializedData.data()), HIPDNN_STATUS_SUCCESS); hipdnn_flatbuffers_sdk::data_objects::GraphT graphT; hipdnn_flatbuffers_sdk::data_objects::GetGraph(serializedData.data())->UnPackTo(&graphT); // Q + K + V + O + dO + Stats + dQ + dK + dV = 9 tensors ASSERT_EQ(graphT.tensors.size(), 9u); // All auto-assigned UIDs should be unique std::unordered_set uids; for(const auto& t : graphT.tensors) { uids.insert(t->uid); } EXPECT_EQ(uids.size(), 9u) << "Tensor UIDs are not unique"; // The SDPA bprop operation should reference the auto-assigned UIDs ASSERT_EQ(graphT.nodes.size(), 1u); auto* sdpa = graphT.nodes[0]->attributes.AsSdpaBackwardAttributes(); ASSERT_NE(sdpa, nullptr); // All tensor UIDs referenced by the node should exist in the graph tensors EXPECT_TRUE(uids.count(sdpa->q_tensor_uid) > 0) << "Q tensor UID " << sdpa->q_tensor_uid << " not found in graph tensors"; EXPECT_TRUE(uids.count(sdpa->k_tensor_uid) > 0) << "K tensor UID " << sdpa->k_tensor_uid << " not found in graph tensors"; EXPECT_TRUE(uids.count(sdpa->v_tensor_uid) > 0) << "V tensor UID " << sdpa->v_tensor_uid << " not found in graph tensors"; EXPECT_TRUE(uids.count(sdpa->o_tensor_uid) > 0) << "O tensor UID " << sdpa->o_tensor_uid << " not found in graph tensors"; EXPECT_TRUE(uids.count(sdpa->do_tensor_uid) > 0) << "dO tensor UID " << sdpa->do_tensor_uid << " not found in graph tensors"; EXPECT_TRUE(uids.count(sdpa->stats_tensor_uid) > 0) << "Stats tensor UID " << sdpa->stats_tensor_uid << " not found in graph tensors"; EXPECT_TRUE(uids.count(sdpa->dq_tensor_uid) > 0) << "dQ tensor UID " << sdpa->dq_tensor_uid << " not found in graph tensors"; EXPECT_TRUE(uids.count(sdpa->dk_tensor_uid) > 0) << "dK tensor UID " << sdpa->dk_tensor_uid << " not found in graph tensors"; EXPECT_TRUE(uids.count(sdpa->dv_tensor_uid) > 0) << "dV tensor UID " << sdpa->dv_tensor_uid << " not found in graph tensors"; // All nine required tensor UIDs referenced by the node should be distinct const std::unordered_set nodeUids = {sdpa->q_tensor_uid, sdpa->k_tensor_uid, sdpa->v_tensor_uid, sdpa->o_tensor_uid, sdpa->do_tensor_uid, sdpa->stats_tensor_uid, sdpa->dq_tensor_uid, sdpa->dk_tensor_uid, sdpa->dv_tensor_uid}; EXPECT_EQ(nodeUids.size(), 9u) << "SDPA bprop node tensor UIDs are not distinct"; } } // namespace