Merged PR 5103319: 8d Update

Required changes for 8D scatter and gather

Related work items: #27678554
This commit is contained in:
Nick Feeney 2020-08-27 21:28:02 +00:00
parent cb5e199a79
commit b5c765c76b
4 changed files with 46 additions and 38 deletions

View file

@ -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));
}
}
}

View file

@ -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;

View file

@ -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];

View file

@ -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();