mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
Merged PR 5103319: 8d Update
Required changes for 8D scatter and gather Related work items: #27678554
This commit is contained in:
parent
cb5e199a79
commit
b5c765c76b
4 changed files with 46 additions and 38 deletions
|
|
@ -133,12 +133,21 @@ namespace Dml
|
|||
));
|
||||
}
|
||||
|
||||
void DmlOperator::Initialize(
|
||||
const MLOperatorKernelCreationContext& kernelInfo,
|
||||
uint32_t minDimensionCount
|
||||
)
|
||||
{
|
||||
Initialize(kernelInfo, std::nullopt, std::nullopt, std::nullopt, std::nullopt, minDimensionCount);
|
||||
}
|
||||
|
||||
void DmlOperator::Initialize(
|
||||
const MLOperatorKernelCreationContext& kernelInfo,
|
||||
const std::optional<const std::vector<std::optional<uint32_t>>>& kernelInputIndices,
|
||||
const std::optional<const std::vector<std::optional<uint32_t>>>& kernelOutputIndices,
|
||||
const std::optional<gsl::span<const uint32_t>> inputShape,
|
||||
const std::optional<gsl::span<const uint32_t>> outputShape
|
||||
const std::optional<gsl::span<const uint32_t>> outputShape,
|
||||
uint32_t minDimensionCount
|
||||
)
|
||||
{
|
||||
if (kernelInputIndices)
|
||||
|
|
@ -179,7 +188,7 @@ namespace Dml
|
|||
TensorAxis::W,
|
||||
TensorAxis::RightAligned,
|
||||
inputShape,
|
||||
NchwDimensionCount));
|
||||
minDimensionCount));
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -200,7 +209,8 @@ namespace Dml
|
|||
TensorAxis::DoNotCoerce,
|
||||
TensorAxis::W,
|
||||
TensorAxis::RightAligned,
|
||||
outputShape));
|
||||
outputShape,
|
||||
minDimensionCount));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -29,12 +29,18 @@ namespace Dml
|
|||
ComPtr<IUnknown> m_persistentResourcePoolingUnk; // Controls when the persistent resource is returned to the pool
|
||||
std::optional<DML_BUFFER_BINDING> m_persistentResourceBinding;
|
||||
|
||||
void Initialize(
|
||||
const MLOperatorKernelCreationContext& kernelInfo,
|
||||
uint32_t minDimensionCount
|
||||
);
|
||||
|
||||
void Initialize(
|
||||
const MLOperatorKernelCreationContext& kernelInfo,
|
||||
const std::optional<const std::vector<std::optional<uint32_t>>>& kernelInputIndices = std::nullopt,
|
||||
const std::optional<const std::vector<std::optional<uint32_t>>>& kernelOutputIndices = std::nullopt,
|
||||
const std::optional<gsl::span<const uint32_t>> inputShape = std::nullopt,
|
||||
const std::optional<gsl::span<const uint32_t>> outputShape = std::nullopt
|
||||
const std::optional<gsl::span<const uint32_t>> outputShape = std::nullopt,
|
||||
uint32_t minDimensionCount = NchwDimensionCount
|
||||
);
|
||||
|
||||
bool AllowHalfPrecisionComputation() const;
|
||||
|
|
|
|||
|
|
@ -16,7 +16,14 @@ public:
|
|||
ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetInputCount() == 2, "Gather expects 2 inputs.");
|
||||
ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetOutputCount() == 1, "Gather expects 1 output.");
|
||||
|
||||
DmlOperator::Initialize(kernelCreationContext);
|
||||
auto tensorShapeDescription = kernelCreationContext.GetTensorShapeDescription();
|
||||
std::vector<DimensionType> dataDimensions = tensorShapeDescription.GetInputTensorShape(0);
|
||||
std::vector<DimensionType> indicesDimensions = tensorShapeDescription.GetInputTensorShape(1);
|
||||
std::vector<DimensionType> outputDimensions = tensorShapeDescription.GetOutputTensorShape(0);
|
||||
|
||||
size_t dimensionCountMax = std::max({dataDimensions.size(), indicesDimensions.size(), outputDimensions.size()});
|
||||
DmlOperator::Initialize(kernelCreationContext, gsl::narrow_cast<uint32_t>(dimensionCountMax));
|
||||
|
||||
DmlOperator::Remap64bitDmlDataTypesTo32bitIfNeeded();
|
||||
|
||||
std::vector<DML_TENSOR_DESC> inputDescs = GetDmlInputDescs();
|
||||
|
|
@ -24,9 +31,6 @@ public:
|
|||
assert(inputDescs.size() == 2);
|
||||
assert(outputDescs.size() == 1);
|
||||
|
||||
auto outputTensorShapeDescription = kernelCreationContext.GetTensorShapeDescription();
|
||||
std::vector<DimensionType> dataDimensions = outputTensorShapeDescription.GetInputTensorShape(0);
|
||||
std::vector<DimensionType> indicesDimensions = outputTensorShapeDescription.GetInputTensorShape(1);
|
||||
uint32_t dmlAxis = GetDmlAdjustedAxis(m_axis, kernelCreationContext, m_inputTensorDescs.front().GetDimensionCount());
|
||||
|
||||
DML_GATHER_OPERATOR_DESC operatorDesc = {};
|
||||
|
|
@ -50,7 +54,14 @@ public:
|
|||
ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetInputCount() == 2, "GatherElements expects 2 inputs.");
|
||||
ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetOutputCount() == 1, "GatherElements expects 1 output.");
|
||||
|
||||
DmlOperator::Initialize(kernelCreationContext);
|
||||
auto tensorShapeDescription = kernelCreationContext.GetTensorShapeDescription();
|
||||
std::vector<DimensionType> dataDimensions = tensorShapeDescription.GetInputTensorShape(0);
|
||||
std::vector<DimensionType> indicesDimensions = tensorShapeDescription.GetInputTensorShape(1);
|
||||
std::vector<DimensionType> outputDimensions = tensorShapeDescription.GetOutputTensorShape(0);
|
||||
|
||||
size_t dimensionCountMax = std::max({dataDimensions.size(), indicesDimensions.size(), outputDimensions.size()});
|
||||
DmlOperator::Initialize(kernelCreationContext, gsl::narrow_cast<uint32_t>(dimensionCountMax));
|
||||
|
||||
DmlOperator::Remap64bitDmlDataTypesTo32bitIfNeeded();
|
||||
|
||||
std::vector<DML_TENSOR_DESC> inputDescs = GetDmlInputDescs();
|
||||
|
|
@ -59,10 +70,6 @@ public:
|
|||
assert(outputDescs.size() == 1);
|
||||
|
||||
int32_t signedOnnxAxis = kernelCreationContext.GetOptionalAttribute<int>(AttrName::Axis, 0);
|
||||
auto outputTensorShapeDescription = kernelCreationContext.GetTensorShapeDescription();
|
||||
std::vector<DimensionType> dataDimensions = outputTensorShapeDescription.GetInputTensorShape(0);
|
||||
std::vector<DimensionType> indicesDimensions = outputTensorShapeDescription.GetInputTensorShape(1);
|
||||
ML_CHECK_VALID_ARGUMENT(dataDimensions.size() <= OperatorHelper::NchwDimensionCount);
|
||||
uint32_t dmlAxis = GetDmlAdjustedAxis(signedOnnxAxis, kernelCreationContext, m_inputTensorDescs.front().GetDimensionCount());
|
||||
|
||||
DML_GATHER_ELEMENTS_OPERATOR_DESC operatorDesc = {};
|
||||
|
|
@ -86,32 +93,21 @@ public:
|
|||
ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetInputCount() == 2, "GatherND expects 2 inputs.");
|
||||
ML_CHECK_VALID_ARGUMENT(kernelCreationContext.GetOutputCount() == 1, "GatherND expects 1 output.");
|
||||
|
||||
DmlOperator::Initialize(kernelCreationContext);
|
||||
auto tensorShapeDescription = kernelCreationContext.GetTensorShapeDescription();
|
||||
std::vector<DimensionType> dataDimensions = tensorShapeDescription.GetInputTensorShape(0);
|
||||
std::vector<DimensionType> indicesDimensions = tensorShapeDescription.GetInputTensorShape(1);
|
||||
std::vector<DimensionType> outputDimensions = tensorShapeDescription.GetOutputTensorShape(0);
|
||||
|
||||
size_t dimensionCountMax = std::max({dataDimensions.size(), indicesDimensions.size(), outputDimensions.size()});
|
||||
DmlOperator::Initialize(kernelCreationContext, gsl::narrow_cast<uint32_t>(dimensionCountMax));
|
||||
|
||||
DmlOperator::Remap64bitDmlDataTypesTo32bitIfNeeded();
|
||||
|
||||
uint32_t maxDimensionCount = std::max({
|
||||
m_inputTensorDescs[0].GetDimensionCount(),
|
||||
m_inputTensorDescs[1].GetDimensionCount(),
|
||||
m_outputTensorDescs[0].GetDimensionCount()
|
||||
});
|
||||
|
||||
// DML expects all tensors to have the same dimension count.
|
||||
// Update the tensor descriptions with new sizes.
|
||||
m_inputTensorDescs[0].SetDimensionCount(maxDimensionCount, TensorAxis::RightAligned);
|
||||
m_inputTensorDescs[1].SetDimensionCount(maxDimensionCount, TensorAxis::RightAligned);
|
||||
m_outputTensorDescs[0].SetDimensionCount(maxDimensionCount, TensorAxis::RightAligned);
|
||||
|
||||
std::vector<DML_TENSOR_DESC> inputDescs = GetDmlInputDescs();
|
||||
std::vector<DML_TENSOR_DESC> outputDescs = GetDmlOutputDescs();
|
||||
assert(inputDescs.size() == 2);
|
||||
assert(outputDescs.size() == 1);
|
||||
|
||||
auto outputTensorShapeDescription = kernelCreationContext.GetTensorShapeDescription();
|
||||
std::vector<DimensionType> dataDimensions = outputTensorShapeDescription.GetInputTensorShape(0);
|
||||
std::vector<DimensionType> indicesDimensions = outputTensorShapeDescription.GetInputTensorShape(1);
|
||||
ML_CHECK_VALID_ARGUMENT(dataDimensions.size() > m_batchCount);
|
||||
ML_CHECK_VALID_ARGUMENT(indicesDimensions.size() > m_batchCount);
|
||||
|
||||
DML_GATHER_ND1_OPERATOR_DESC operatorDesc = {};
|
||||
operatorDesc.InputTensor = &inputDescs[0];
|
||||
operatorDesc.IndicesTensor = &inputDescs[1];
|
||||
|
|
|
|||
|
|
@ -23,7 +23,6 @@ public:
|
|||
ML_CHECK_VALID_ARGUMENT(dataDimensions == outputDimensions);
|
||||
ML_CHECK_VALID_ARGUMENT(indicesDimensions == updatesDimensions);
|
||||
ML_CHECK_VALID_ARGUMENT(dataDimensions.size() == indicesDimensions.size());
|
||||
ML_CHECK_VALID_ARGUMENT(dataDimensions.size() <= OperatorHelper::NchwDimensionCount);
|
||||
|
||||
// When the indices tensor is empty, Scatter is basically Identity. But since DML doesn't support empty or null
|
||||
// tensors, we have to special-case it outside of DML.
|
||||
|
|
@ -89,12 +88,9 @@ public:
|
|||
std::vector<DimensionType> updatesDimensions = tensorShapeDescription.GetInputTensorShape(2);
|
||||
std::vector<DimensionType> outputDimensions = tensorShapeDescription.GetOutputTensorShape(0);
|
||||
ML_CHECK_VALID_ARGUMENT(dataDimensions == outputDimensions);
|
||||
ML_CHECK_VALID_ARGUMENT(dataDimensions.size() <= OperatorHelper::NchwDimensionCount);
|
||||
ML_CHECK_VALID_ARGUMENT(indicesDimensions.size() <= OperatorHelper::NchwDimensionCount);
|
||||
ML_CHECK_VALID_ARGUMENT(updatesDimensions.size() <= OperatorHelper::NchwDimensionCount);
|
||||
ML_CHECK_VALID_ARGUMENT(outputDimensions.size() <= OperatorHelper::NchwDimensionCount);
|
||||
|
||||
DmlOperator::Initialize(kernelCreationContext);
|
||||
size_t dimensionCountMax = std::max({dataDimensions.size(), updatesDimensions.size(), indicesDimensions.size(), outputDimensions.size()});
|
||||
DmlOperator::Initialize(kernelCreationContext, gsl::narrow_cast<uint32_t>(dimensionCountMax));
|
||||
DmlOperator::Remap64bitDmlDataTypesTo32bitIfNeeded();
|
||||
|
||||
std::vector<DML_TENSOR_DESC> inputDescs = GetDmlInputDescs();
|
||||
|
|
|
|||
Loading…
Reference in a new issue