Support exteranl weights in DML execution provider (#13740)

### Description
This enables support for external weights in the DML execution provider
when its graph optimization logic is reached.

### Motivation and Context
External weighs are encountered after optimization is applied to
transformer models.
This commit is contained in:
Jeff Bloomfield 2022-11-29 15:47:16 -08:00 committed by GitHub
parent ce0025d3f2
commit 571dc5a1f1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 48 additions and 35 deletions

View file

@ -163,21 +163,29 @@ namespace DmlGraphFusionHelper
if (iter != initializerNameToInitializerMap.end())
{
std::byte* tensorPtr = nullptr;
size_t tensorByteSize = 0;
size_t tensorByteSize = 0;
std::vector<uint8_t> unpackedExternalTensor;
std::unique_ptr<std::byte[]> unpackedTensor;
//auto& initializer = iter->second;
auto* initializer = iter->second.first;
// The tensor may be stored as raw data or in typed fields.
if (initializer->has_raw_data())
if (initializer->data_location() == onnx::TensorProto_DataLocation_EXTERNAL)
{
THROW_IF_NOT_OK(onnxruntime::utils::UnpackInitializerData(*initializer, graph.ModelPath(), unpackedExternalTensor));
tensorPtr = reinterpret_cast<std::byte*>(unpackedExternalTensor.data());
tensorByteSize = unpackedExternalTensor.size();
}
else if (initializer->has_raw_data())
{
tensorPtr = (std::byte*)(initializer->raw_data().c_str());
tensorByteSize = initializer->raw_data().size();
}
else
{
std::tie(unpackedTensor, tensorByteSize) = Windows::AI::MachineLearning::Adapter::UnpackTensor(*initializer);
std::tie(unpackedTensor, tensorByteSize) = Windows::AI::MachineLearning::Adapter::UnpackTensor(*initializer, graph.ModelPath());
tensorPtr = unpackedTensor.get();
// Free the initializer if this is the last usage of it.

View file

@ -90,7 +90,9 @@ namespace Dml::GraphDescBuilder
reuseCommandList = true;
}
auto constantCpuGraphInputGetter = [&isInitializerTransferable](const std::string& argName)
auto modelPath = graph.ModelPath();
auto constantCpuGraphInputGetter = [&isInitializerTransferable, &modelPath](const std::string& argName)
{
ComPtr<OnnxTensorWrapper> tensorWrapper;
@ -98,7 +100,7 @@ namespace Dml::GraphDescBuilder
if (iter != isInitializerTransferable.end())
{
// Using const_cast here is simpler than making surrounding code const correct.
tensorWrapper = wil::MakeOrThrow<OnnxTensorWrapper>(const_cast<ONNX_NAMESPACE::TensorProto*>(iter->second.first));
tensorWrapper = wil::MakeOrThrow<OnnxTensorWrapper>(const_cast<ONNX_NAMESPACE::TensorProto*>(iter->second.first), modelPath);
}
return tensorWrapper;

View file

@ -705,18 +705,20 @@ namespace Windows::AI::MachineLearning::Adapter
*tensor = nullptr;
// Read the tensor if present, and wrap it in a IMLOperatorTensor.
const onnx::AttributeProto* attributeProto = m_impl->TryGetAttribute(std::string(name));
if (attributeProto)
// Read the tensor if present, and wrap it in a IMLOperatorTensor.
const onnx::AttributeProto* attributeProto = m_impl->TryGetAttribute(std::string(name));
if (attributeProto)
{
if (attributeProto->has_t())
{
if (attributeProto->has_t())
{
const onnx::TensorProto* tensorProto = &attributeProto->t();
Microsoft::WRL::ComPtr<IMLOperatorTensor> tensorWrapper = wil::MakeOrThrow<OnnxTensorWrapper>(const_cast<onnx::TensorProto*>(tensorProto));
*tensor = tensorWrapper.Detach();
return S_OK;
}
const onnx::TensorProto* tensorProto = &attributeProto->t();
// An empty path is used as external weights are not currently supported in this case
Microsoft::WRL::ComPtr<IMLOperatorTensor> tensorWrapper = wil::MakeOrThrow<OnnxTensorWrapper>(const_cast<onnx::TensorProto*>(tensorProto), onnxruntime::Path());
*tensor = tensorWrapper.Detach();
return S_OK;
}
}
return E_INVALIDARG; // The argument has no valid matching attribute.
}
@ -1186,7 +1188,7 @@ namespace Windows::AI::MachineLearning::Adapter
ORT_CATCH_RETURN
}
OnnxTensorWrapper::OnnxTensorWrapper(onnx::TensorProto* impl) : m_impl(impl)
OnnxTensorWrapper::OnnxTensorWrapper(onnx::TensorProto* impl, const onnxruntime::Path& modelPath) : m_impl(impl)
{
// The tensor may be stored as raw data or in typed fields.
if (impl->has_raw_data())
@ -1196,7 +1198,7 @@ namespace Windows::AI::MachineLearning::Adapter
}
else
{
std::tie(m_unpackedTensor, m_tensorByteSize) = UnpackTensor(*impl);
std::tie(m_unpackedTensor, m_tensorByteSize) = UnpackTensor(*impl, modelPath);
m_dataPtr = m_unpackedTensor.get();
}
}
@ -2111,8 +2113,9 @@ namespace Windows::AI::MachineLearning::Adapter
MLOperatorTensorGetter mlOperatorTensorGetter = MLOperatorTensorGetter(
[ctx](uint32_t index)
{
// An empty path is used as external weights are not currently supported in this case
Microsoft::WRL::ComPtr<IMLOperatorTensor> tensorWrapper = wil::MakeOrThrow<OnnxTensorWrapper>(
const_cast<onnx::TensorProto*>(ctx->getInputData(index)));
const_cast<onnx::TensorProto*>(ctx->getInputData(index)), onnxruntime::Path());
return tensorWrapper;
}
);
@ -2303,25 +2306,25 @@ namespace Windows::AI::MachineLearning::Adapter
return false;
}
std::tuple<std::unique_ptr<std::byte[]>, size_t> UnpackTensor(const onnx::TensorProto& initializer)
std::tuple<std::unique_ptr<std::byte[]>, size_t> UnpackTensor(
const onnx::TensorProto& initializer,
const onnxruntime::Path& modelPath)
{
std::unique_ptr<std::byte[]> unpackedTensor;
size_t tensorByteSize = 0;
#define CASE_PROTO(X, Y, Z) \
case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_##X: \
{ \
size_t elementCount = initializer.##Z(); \
tensorByteSize = elementCount * sizeof(Y); \
unpackedTensor.reset(new std::byte[tensorByteSize]); \
ORT_THROW_HR_IF(E_FAIL, !onnxruntime::utils::UnpackTensor( \
initializer, \
initializer.has_raw_data() ? initializer.raw_data().data() : nullptr, \
initializer.has_raw_data() ? initializer.raw_data().size() : 0, \
reinterpret_cast<Y*>(unpackedTensor.get()), elementCount) \
.IsOK()); \
break; \
}
#define CASE_PROTO(X, Y, Z) \
case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_##X: { \
size_t elementCount = initializer.##Z(); \
tensorByteSize = elementCount * sizeof(Y); \
unpackedTensor.reset(new std::byte[tensorByteSize]); \
ORT_THROW_HR_IF(E_FAIL, !onnxruntime::utils::UnpackTensor( \
initializer, \
modelPath, \
reinterpret_cast<Y*>(unpackedTensor.get()), elementCount) \
.IsOK()); \
break; \
}
switch (initializer.data_type())
{
CASE_PROTO(FLOAT, float, float_data_size);

View file

@ -296,7 +296,7 @@ class OnnxTensorWrapper : public WRL::Base<IMLOperatorTensor>, public Closable
public:
OnnxTensorWrapper() = default;
OnnxTensorWrapper(onnx::TensorProto* impl);
OnnxTensorWrapper(onnx::TensorProto* impl, const onnxruntime::Path& modelPath);
uint32_t STDMETHODCALLTYPE GetDimensionCount() const noexcept override;
@ -648,5 +648,5 @@ bool TryGetStaticInputShapes(const onnxruntime::Node& node, EdgeShapes& inputSha
bool TryGetStaticOutputShapes(const onnxruntime::Node& node, EdgeShapes& outputShapes);
bool ContainsEmptyDimensions(const EdgeShapes& shapes, gsl::span<const uint32_t> ignoredShapeIndices = gsl::span<const uint32_t>());
std::tuple<std::unique_ptr<std::byte[]>, size_t> UnpackTensor(const onnx::TensorProto& initializer);
std::tuple<std::unique_ptr<std::byte[]>, size_t> UnpackTensor(const onnx::TensorProto& initializer, const onnxruntime::Path& modelPath);
} // namespace Windows::AI::MachineLearning::Adapter