derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
280 lines
12 KiB
C++
280 lines
12 KiB
C++
// Copyright © Advanced Micro Devices, Inc., or its affiliates.
|
|
// SPDX-License-Identifier: MIT
|
|
|
|
#pragma once
|
|
|
|
#include <hipdnn_frontend/Types.hpp>
|
|
#include <hipdnn_frontend/attributes/SdpaBackwardAttributes.hpp>
|
|
#include <hipdnn_frontend/detail/DescriptorUnpackHelpers.hpp>
|
|
#include <memory>
|
|
#include <unordered_map>
|
|
|
|
namespace hipdnn_frontend::detail
|
|
{
|
|
|
|
/// Unpacks an SDPA backward operation descriptor and populates
|
|
/// SdpaBackwardAttributes with tensors (using tensorMap for sharing)
|
|
/// and scalar/enum parameters.
|
|
[[nodiscard]] inline Error unpackSdpaBwdOperation(
|
|
hipdnnBackendDescriptor_t opDesc,
|
|
std::unordered_map<int64_t, std::shared_ptr<graph::TensorAttributes>>& tensorMap,
|
|
graph::SdpaBackwardAttributes& attributes)
|
|
{
|
|
// Required input tensors
|
|
std::shared_ptr<graph::TensorAttributes> qTensor;
|
|
HIPDNN_CHECK_ERROR(unpackAndRegisterTensor(
|
|
opDesc, HIPDNN_ATTR_OPERATION_SDPA_BWD_Q_EXT, tensorMap, qTensor, "SDPA bprop Q"));
|
|
attributes.set_q(qTensor);
|
|
|
|
std::shared_ptr<graph::TensorAttributes> kTensor;
|
|
HIPDNN_CHECK_ERROR(unpackAndRegisterTensor(
|
|
opDesc, HIPDNN_ATTR_OPERATION_SDPA_BWD_K_EXT, tensorMap, kTensor, "SDPA bprop K"));
|
|
attributes.set_k(kTensor);
|
|
|
|
std::shared_ptr<graph::TensorAttributes> vTensor;
|
|
HIPDNN_CHECK_ERROR(unpackAndRegisterTensor(
|
|
opDesc, HIPDNN_ATTR_OPERATION_SDPA_BWD_V_EXT, tensorMap, vTensor, "SDPA bprop V"));
|
|
attributes.set_v(vTensor);
|
|
|
|
std::shared_ptr<graph::TensorAttributes> oTensor;
|
|
HIPDNN_CHECK_ERROR(unpackAndRegisterTensor(
|
|
opDesc, HIPDNN_ATTR_OPERATION_SDPA_BWD_O_EXT, tensorMap, oTensor, "SDPA bprop O"));
|
|
attributes.set_o(oTensor);
|
|
|
|
std::shared_ptr<graph::TensorAttributes> doTensor;
|
|
HIPDNN_CHECK_ERROR(unpackAndRegisterTensor(
|
|
opDesc, HIPDNN_ATTR_OPERATION_SDPA_BWD_DO_EXT, tensorMap, doTensor, "SDPA bprop dO"));
|
|
attributes.set_do(doTensor);
|
|
|
|
std::shared_ptr<graph::TensorAttributes> statsTensor;
|
|
HIPDNN_CHECK_ERROR(unpackAndRegisterTensor(opDesc,
|
|
HIPDNN_ATTR_OPERATION_SDPA_BWD_STATS_EXT,
|
|
tensorMap,
|
|
statsTensor,
|
|
"SDPA bprop Stats"));
|
|
attributes.set_stats(statsTensor);
|
|
|
|
// Required output tensors
|
|
std::shared_ptr<graph::TensorAttributes> dqTensor;
|
|
HIPDNN_CHECK_ERROR(unpackAndRegisterTensor(
|
|
opDesc, HIPDNN_ATTR_OPERATION_SDPA_BWD_DQ_EXT, tensorMap, dqTensor, "SDPA bprop dQ"));
|
|
attributes.set_dq(dqTensor);
|
|
|
|
std::shared_ptr<graph::TensorAttributes> dkTensor;
|
|
HIPDNN_CHECK_ERROR(unpackAndRegisterTensor(
|
|
opDesc, HIPDNN_ATTR_OPERATION_SDPA_BWD_DK_EXT, tensorMap, dkTensor, "SDPA bprop dK"));
|
|
attributes.set_dk(dkTensor);
|
|
|
|
std::shared_ptr<graph::TensorAttributes> dvTensor;
|
|
HIPDNN_CHECK_ERROR(unpackAndRegisterTensor(
|
|
opDesc, HIPDNN_ATTR_OPERATION_SDPA_BWD_DV_EXT, tensorMap, dvTensor, "SDPA bprop dV"));
|
|
attributes.set_dv(dvTensor);
|
|
|
|
// Optional input tensors
|
|
std::shared_ptr<graph::TensorAttributes> scaleTensor;
|
|
HIPDNN_CHECK_ERROR(unpackOptionalTensor(opDesc,
|
|
HIPDNN_ATTR_OPERATION_SDPA_BWD_SCALE_EXT,
|
|
tensorMap,
|
|
scaleTensor,
|
|
"SDPA bprop SCALE"));
|
|
if(scaleTensor)
|
|
{
|
|
attributes.set_attn_scale(scaleTensor);
|
|
}
|
|
|
|
std::shared_ptr<graph::TensorAttributes> attnMaskTensor;
|
|
HIPDNN_CHECK_ERROR(unpackOptionalTensor(opDesc,
|
|
HIPDNN_ATTR_OPERATION_SDPA_BWD_ATTN_MASK_EXT,
|
|
tensorMap,
|
|
attnMaskTensor,
|
|
"SDPA bprop ATTN_MASK"));
|
|
if(attnMaskTensor)
|
|
{
|
|
// The attention mask is stored via set_bias() in SdpaBackwardAttributes
|
|
// (same naming convention as the forward pass).
|
|
attributes.set_bias(attnMaskTensor);
|
|
}
|
|
|
|
std::shared_ptr<graph::TensorAttributes> seqLenQTensor;
|
|
HIPDNN_CHECK_ERROR(unpackOptionalTensor(opDesc,
|
|
HIPDNN_ATTR_OPERATION_SDPA_BWD_SEQ_LEN_Q_EXT,
|
|
tensorMap,
|
|
seqLenQTensor,
|
|
"SDPA bprop SEQ_LEN_Q"));
|
|
if(seqLenQTensor)
|
|
{
|
|
attributes.set_seq_len_q(seqLenQTensor);
|
|
}
|
|
|
|
std::shared_ptr<graph::TensorAttributes> seqLenKvTensor;
|
|
HIPDNN_CHECK_ERROR(unpackOptionalTensor(opDesc,
|
|
HIPDNN_ATTR_OPERATION_SDPA_BWD_SEQ_LEN_KV_EXT,
|
|
tensorMap,
|
|
seqLenKvTensor,
|
|
"SDPA bprop SEQ_LEN_KV"));
|
|
if(seqLenKvTensor)
|
|
{
|
|
attributes.set_seq_len_kv(seqLenKvTensor);
|
|
}
|
|
|
|
std::shared_ptr<graph::TensorAttributes> seedTensor;
|
|
HIPDNN_CHECK_ERROR(unpackOptionalTensor(
|
|
opDesc, HIPDNN_ATTR_OPERATION_SDPA_BWD_SEED_EXT, tensorMap, seedTensor, "SDPA bprop SEED"));
|
|
if(seedTensor)
|
|
{
|
|
attributes.set_seed(seedTensor);
|
|
}
|
|
|
|
std::shared_ptr<graph::TensorAttributes> offsetTensor;
|
|
HIPDNN_CHECK_ERROR(unpackOptionalTensor(opDesc,
|
|
HIPDNN_ATTR_OPERATION_SDPA_BWD_OFFSET_EXT,
|
|
tensorMap,
|
|
offsetTensor,
|
|
"SDPA bprop OFFSET"));
|
|
if(offsetTensor)
|
|
{
|
|
attributes.set_offset(offsetTensor);
|
|
}
|
|
|
|
std::shared_ptr<graph::TensorAttributes> dropoutMaskTensor;
|
|
HIPDNN_CHECK_ERROR(unpackOptionalTensor(opDesc,
|
|
HIPDNN_ATTR_OPERATION_SDPA_BWD_DROPOUT_MASK_EXT,
|
|
tensorMap,
|
|
dropoutMaskTensor,
|
|
"SDPA bprop DROPOUT_MASK"));
|
|
if(dropoutMaskTensor)
|
|
{
|
|
attributes.set_dropout_mask(dropoutMaskTensor);
|
|
}
|
|
|
|
std::shared_ptr<graph::TensorAttributes> dropoutScaleTensor;
|
|
HIPDNN_CHECK_ERROR(unpackOptionalTensor(opDesc,
|
|
HIPDNN_ATTR_OPERATION_SDPA_BWD_DROPOUT_SCALE_EXT,
|
|
tensorMap,
|
|
dropoutScaleTensor,
|
|
"SDPA bprop DROPOUT_SCALE"));
|
|
if(dropoutScaleTensor)
|
|
{
|
|
attributes.set_dropout_scale(dropoutScaleTensor);
|
|
}
|
|
|
|
std::shared_ptr<graph::TensorAttributes> dropoutScaleInvTensor;
|
|
HIPDNN_CHECK_ERROR(unpackOptionalTensor(opDesc,
|
|
HIPDNN_ATTR_OPERATION_SDPA_BWD_DROPOUT_SCALE_INV_EXT,
|
|
tensorMap,
|
|
dropoutScaleInvTensor,
|
|
"SDPA bprop DROPOUT_SCALE_INV"));
|
|
if(dropoutScaleInvTensor)
|
|
{
|
|
attributes.set_dropout_scale_inv(dropoutScaleInvTensor);
|
|
}
|
|
|
|
// Optional output tensor
|
|
std::shared_ptr<graph::TensorAttributes> dbiasTensor;
|
|
HIPDNN_CHECK_ERROR(unpackOptionalTensor(opDesc,
|
|
HIPDNN_ATTR_OPERATION_SDPA_BWD_DBIAS_EXT,
|
|
tensorMap,
|
|
dbiasTensor,
|
|
"SDPA bprop DBIAS"));
|
|
if(dbiasTensor)
|
|
{
|
|
attributes.set_dbias(dbiasTensor);
|
|
}
|
|
|
|
// Boolean scalars — plain bool members, so unpack into optional and apply default
|
|
{
|
|
std::optional<bool> opt;
|
|
HIPDNN_CHECK_ERROR(getDescriptorAttrOptionalScalar(opDesc,
|
|
HIPDNN_ATTR_SDPA_BWD_ALIBI_MASK_EXT,
|
|
HIPDNN_TYPE_BOOLEAN,
|
|
opt,
|
|
"SDPA bprop alibi_mask"));
|
|
attributes.alibi_mask = opt.value_or(false);
|
|
}
|
|
{
|
|
std::optional<bool> opt;
|
|
HIPDNN_CHECK_ERROR(getDescriptorAttrOptionalScalar(opDesc,
|
|
HIPDNN_ATTR_SDPA_BWD_PADDING_MASK_EXT,
|
|
HIPDNN_TYPE_BOOLEAN,
|
|
opt,
|
|
"SDPA bprop padding_mask"));
|
|
attributes.padding_mask = opt.value_or(false);
|
|
}
|
|
{
|
|
std::optional<bool> opt;
|
|
HIPDNN_CHECK_ERROR(getDescriptorAttrOptionalScalar(opDesc,
|
|
HIPDNN_ATTR_SDPA_BWD_CAUSAL_MASK_EXT,
|
|
HIPDNN_TYPE_BOOLEAN,
|
|
opt,
|
|
"SDPA bprop causal_mask"));
|
|
attributes.causal_mask = opt.value_or(false);
|
|
}
|
|
{
|
|
std::optional<bool> opt;
|
|
HIPDNN_CHECK_ERROR(
|
|
getDescriptorAttrOptionalScalar(opDesc,
|
|
HIPDNN_ATTR_SDPA_BWD_CAUSAL_MASK_BOTTOM_RIGHT_EXT,
|
|
HIPDNN_TYPE_BOOLEAN,
|
|
opt,
|
|
"SDPA bprop causal_mask_bottom_right"));
|
|
attributes.causal_mask_bottom_right = opt.value_or(false);
|
|
}
|
|
|
|
// Optional float scalars
|
|
HIPDNN_CHECK_ERROR(getDescriptorAttrOptionalScalar(opDesc,
|
|
HIPDNN_ATTR_SDPA_BWD_DROPOUT_PROBABILITY_EXT,
|
|
HIPDNN_TYPE_FLOAT,
|
|
attributes.dropout_probability,
|
|
"SDPA bprop dropout_probability"));
|
|
HIPDNN_CHECK_ERROR(getDescriptorAttrOptionalScalar(opDesc,
|
|
HIPDNN_ATTR_SDPA_BWD_ATTN_SCALE_VALUE_EXT,
|
|
HIPDNN_TYPE_FLOAT,
|
|
attributes.attn_scale_value,
|
|
"SDPA bprop attn_scale_value"));
|
|
|
|
// Optional int64 scalars
|
|
HIPDNN_CHECK_ERROR(getDescriptorAttrOptionalScalar(opDesc,
|
|
HIPDNN_ATTR_SDPA_BWD_LEFT_BOUND_EXT,
|
|
HIPDNN_TYPE_INT64,
|
|
attributes.left_bound,
|
|
"SDPA bprop left_bound"));
|
|
HIPDNN_CHECK_ERROR(getDescriptorAttrOptionalScalar(opDesc,
|
|
HIPDNN_ATTR_SDPA_BWD_RIGHT_BOUND_EXT,
|
|
HIPDNN_TYPE_INT64,
|
|
attributes.right_bound,
|
|
"SDPA bprop right_bound"));
|
|
|
|
// Diagonal alignment
|
|
hipdnnDiagonalAlignment_t diagAlign{};
|
|
HIPDNN_CHECK_ERROR(getDescriptorAttrScalar(opDesc,
|
|
HIPDNN_ATTR_SDPA_BWD_DIAGONAL_ALIGNMENT_EXT,
|
|
HIPDNN_TYPE_DIAGONAL_ALIGNMENT_EXT,
|
|
diagAlign,
|
|
"SDPA bprop diagonal_alignment"));
|
|
auto [alignVal, alignErr] = fromHipdnnDiagonalAlignment(diagAlign);
|
|
if(alignErr.is_bad())
|
|
{
|
|
return alignErr;
|
|
}
|
|
attributes.set_diagonal_alignment(alignVal);
|
|
|
|
// Compute data type
|
|
auto [dt, dtErr]
|
|
= unpackGraphDataType(opDesc, HIPDNN_ATTR_SDPA_BWD_COMP_TYPE_EXT, "SDPA bprop math prec");
|
|
if(dtErr.is_bad())
|
|
{
|
|
return dtErr;
|
|
}
|
|
attributes.set_compute_data_type(dt);
|
|
|
|
// Operation name
|
|
std::string opName;
|
|
HIPDNN_CHECK_ERROR(
|
|
getDescriptorAttrString(opDesc, HIPDNN_ATTR_OPERATION_NAME_EXT, opName, "operation name"));
|
|
attributes.set_name(opName);
|
|
|
|
return {};
|
|
}
|
|
|
|
} // namespace hipdnn_frontend::detail
|