Skip to content

Commit

Permalink
DFT-20 version converter (#5613)
Browse files Browse the repository at this point in the history
Implement version converter to complete DFT-20.

Requires #5514 

Reference: #5514 (comment)

---------

Signed-off-by: Justin Chu <justinchu@microsoft.com>
Signed-off-by: Ganesan Ramalingam <grama@microsoft.com>
Co-authored-by: Ganesan Ramalingam <grama@microsoft.com>
  • Loading branch information
justinchuby and gramalingam committed Sep 27, 2023
1 parent b5111b8 commit e11dacf
Show file tree
Hide file tree
Showing 6 changed files with 265 additions and 17 deletions.
48 changes: 34 additions & 14 deletions onnx/test/version_converter/automatic_conversion_test_base.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
# Copyright (c) ONNX Project Contributors

# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations

import string
import unittest
from typing import Any, Dict, List, Optional, Sequence, Union, cast
from typing import Any, List, Sequence, cast

import onnx
from onnx import TensorProto, ValueInfoProto, helper, shape_inference, version_converter
Expand All @@ -12,16 +14,39 @@


class TestAutomaticConversion(unittest.TestCase):
def _test_model_conversion(
self, to_opset: int, model: str | onnx.ModelProto
) -> None:
if isinstance(model, str):
model = onnx.parser.parse_model(model)
onnx.checker.check_model(model)
shape_inference.infer_shapes(model, strict_mode=True)

converted = version_converter.convert_version(model, to_opset)
onnx.checker.check_model(converted)
shape_inference.infer_shapes(converted, strict_mode=True)

def _test_model_conversion_fails(
self, to_opset: int, model: str | onnx.ModelProto
) -> None:
if isinstance(model, str):
model = onnx.parser.parse_model(model)
onnx.checker.check_model(model)
shape_inference.infer_shapes(model, strict_mode=True)

with self.assertRaises(RuntimeError):
version_converter.convert_version(model, to_opset)

def _test_op_conversion(
self,
op: str,
from_opset: int,
input_shapes: Sequence[Union[Sequence[Optional[int]], str]] = ((3, 4, 5),),
output_shapes: Sequence[Sequence[Optional[int]]] = ((3, 4, 5),),
input_types: Optional[Sequence[Any]] = None,
output_types: Optional[Sequence[Any]] = None,
input_shapes: Sequence[Sequence[int | None] | str] = ((3, 4, 5),),
output_shapes: Sequence[Sequence[int | None]] = ((3, 4, 5),),
input_types: Sequence[Any] | None = None,
output_types: Sequence[Any] | None = None,
initializer: Sequence[Any] = (),
attrs: Optional[Dict[str, Any]] = None,
attrs: dict[str, Any] | None = None,
seq_inputs: Sequence[int] = (),
seq_outputs: Sequence[int] = (),
optional_inputs: Sequence[int] = (),
Expand Down Expand Up @@ -67,7 +92,7 @@ def _test_op_conversion(
List[List[int]],
[[0] if isinstance(shape, str) else shape for shape in input_shapes],
)
inputs: List[ValueInfoProto] = []
inputs: list[ValueInfoProto] = []
for name, ttype, shape, is_seq, is_opt in zip(
input_names, input_types, input_shapes_cast, is_sequence, is_optional
):
Expand Down Expand Up @@ -95,7 +120,7 @@ def _test_op_conversion(
List[List[int]],
[[0] if isinstance(shape, str) else shape for shape in output_shapes],
)
outputs: List[ValueInfoProto] = []
outputs: list[ValueInfoProto] = []
for name, ttype, shape, is_seq, is_opt in zip(
output_names, output_types, output_shapes_cast, is_sequence, is_optional
):
Expand All @@ -117,9 +142,4 @@ def _test_op_conversion(
producer_name="test",
opset_imports=[helper.make_opsetid("", start_opset)],
)
onnx.checker.check_model(original)
shape_inference.infer_shapes(original, strict_mode=True)

converted = version_converter.convert_version(original, end_opset)
onnx.checker.check_model(converted)
shape_inference.infer_shapes(converted, strict_mode=True)
self._test_model_conversion(end_opset, original)
50 changes: 50 additions & 0 deletions onnx/test/version_converter/automatic_downgrade_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,56 @@ def test_reduce_ops(self, op) -> None:
initializer=[axes],
)

def test_dft20_no_axis(self) -> None:
self._test_model_conversion(
to_opset=19,
model="""
<ir_version: 9, opset_import: [ "" : 20]>
dft_no_axis (float[N, M, 1] x) => (float[N, M, 2] y)
{
y = DFT (x)
}
""",
)

def test_dft20_initializer_axis(self) -> None:
self._test_model_conversion(
to_opset=19,
model="""
<ir_version: 9, opset_import: [ "" : 20]>
dft_no_axis (float[N, M, 1] x, int64 dft_length) => (float[N, K, 2] y)
<int64 axis = {1}>
{
y = DFT (x, dft_length, axis)
}
""",
)

def test_dft20_constant_axis(self) -> None:
self._test_model_conversion(
to_opset=19,
model="""
<ir_version: 9, opset_import: [ "" : 20]>
dft_no_axis (float[N, M, 1] x, int64 dft_length) => (float[N, K, 2] y)
{
axis = Constant <value = int64{1}>()
y = DFT (x, dft_length, axis)
}
""",
)

def test_dft20_unknown_axis(self) -> None:
self._test_model_conversion_fails(
to_opset=19,
model="""
<ir_version: 9, opset_import: [ "" : 20]>
dft_no_axis (float[N, M, 1] x, int64 dft_length, int64 axis) => (float[P, K, 2] y)
{
y = DFT (x, dft_length, axis)
}
""",
)


if __name__ == "__main__":
unittest.main()
5 changes: 3 additions & 2 deletions onnx/test/version_converter/automatic_upgrade_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@

import automatic_conversion_test_base
import numpy as np
import pytest

import onnx
from onnx import TensorProto, helper
Expand Down Expand Up @@ -1442,7 +1441,6 @@ def test_HannWindow(self) -> None:
def test_HammingWindow(self) -> None:
self._test_window_function("HammingWindow")

@pytest.mark.xfail(reason="FIXME(#5613): Implement version converters for DFT")
def test_DFT(self) -> None:
self._test_op_upgrade("DFT", 17, [[2, 16, 1], []], [[2, 16, 2]])
self._test_op_upgrade("DFT", 17, [[2, 16, 2], []], [[2, 16, 2]])
Expand All @@ -1458,6 +1456,9 @@ def test_DFT(self) -> None:
self._test_op_upgrade(
"DFT", 17, [[2, 16, 2], []], [[2, 16, 2]], attrs={"inverse": 1}
)
self._test_op_upgrade(
"DFT", 17, [[2, 16, 2], []], [[2, 16, 2]], attrs={"inverse": 1, "axis": 0}
)

def _test_short_time_fourier_transform(self, operator_name: str) -> None:
# Real
Expand Down
74 changes: 74 additions & 0 deletions onnx/version_converter/adapters/axis_attribute_to_input.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
// Copyright (c) ONNX Project Contributors

/*
* SPDX-License-Identifier: Apache-2.0
*/

#pragma once

#include <memory>
#include <string>
#include <vector>

#include "onnx/version_converter/adapters/adapter.h"

namespace ONNX_NAMESPACE {
namespace version_conversion {

class AxisAttributeToInput : public Adapter {
public:
AxisAttributeToInput(
const std::string& op_name,
const OpSetID& initial,
const OpSetID& target,
size_t axis_index,
int64_t default_axis)
: Adapter(op_name, initial, target), axis_index(axis_index), default_axis(default_axis) {}

Node* adapt(std::shared_ptr<Graph> graph, Node* node) const override {
if (node->hasAttribute(kaxis)) {
AttrToInput(graph, node, node->i(kaxis), this->axis_index);
node->removeAttribute(kaxis);
return node;
}

// Fill in the default value for axis
AttrToInput(graph, node, default_axis, this->axis_index);
return node;
}

private:
size_t axis_index;
int64_t default_axis;

void AttrToInput(std::shared_ptr<Graph> graph, Node* node, int64_t axis, size_t axis_index) const {
const ArrayRef<Value*>& inputs = node->inputs();

// Add the optional inputs if they don't exist
for (size_t i = inputs.size(); i < axis_index; ++i) {
Node* empty_input = graph->create(kUndefined);
empty_input->insertBefore(node);
node->addInput(empty_input->output());
}

// Add the axis input
Node* constant = CreateAxisInput(graph, node, axis);
node->addInput(constant->output());
}

Node* CreateAxisInput(std::shared_ptr<Graph> graph, Node* node, int64_t axis) const {
Tensor t;
t.elem_type() = TensorProto_DataType_INT64;
t.sizes() = std::vector<int64_t>{};
auto& data = t.int64s();
data.emplace_back(axis);

Node* constant = graph->create(kConstant);
constant->insertBefore(node);
constant->t_(kvalue, t);
return constant;
}
};

} // namespace version_conversion
} // namespace ONNX_NAMESPACE
99 changes: 99 additions & 0 deletions onnx/version_converter/adapters/axis_input_to_attribute.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,99 @@
// Copyright (c) ONNX Project Contributors

/*
* SPDX-License-Identifier: Apache-2.0
*/

#pragma once

#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "onnx/version_converter/adapters/adapter.h"

namespace ONNX_NAMESPACE {
namespace version_conversion {
class AxisInputToAttribute : public Adapter {
public:
explicit AxisInputToAttribute(
const std::string& op_name,
const OpSetID& initial,
const OpSetID& target,
size_t axis_index,
int64_t default_axis)
: Adapter(op_name, initial, target), axis_index(axis_index), default_axis(default_axis) {}

Node* adapt(std::shared_ptr<Graph> graph, Node* node) const override {
if (!HasAxisInput(node)) {
node->i_(kaxis, this->default_axis);
return EnsureAndReturnNode(node);
}

const ArrayRef<Value*>& inputs = node->inputs();
Value* axis_val = inputs[this->axis_index];
Node* axis_node = axis_val->node();

if (axis_node->kind() == kConstant) {
HandleConstantNode(node, axis_node, axis_val);
return EnsureAndReturnNode(node);
}

if (graph->is_constant_initializer(axis_val)) {
HandleInitializerNode(graph, node, axis_val);
return EnsureAndReturnNode(node);
}

ONNX_ASSERTM(false, "Axis input must be a constant or initializer for promotion to attribute.");
}

private:
size_t axis_index;
int64_t default_axis;

bool HasAxisInput(const Node* node) const {
const ArrayRef<const Value*>& inputs = node->inputs();
return inputs.size() > this->axis_index && inputs[this->axis_index]->node()->kind() != kUndefined;
}

void HandleConstantNode(Node* node, Node* axis_node, Value* axis_val) const {
const std::vector<int64_t>& int64s = axis_node->t(kvalue).int64s();
if (int64s.empty()) {
std::string raw_data = axis_node->t(kvalue).raw();
ONNX_ASSERTM(
raw_data.size() != 0 && raw_data.size() % 8 == 0,
"Raw Data must be non-empty and size must be a multiple of 8");
const int64_t* raw = reinterpret_cast<const int64_t*>(raw_data.c_str());
node->i_(kaxis, raw[0]);
} else {
node->i_(kaxis, int64s.at(0));
}
node->removeInput(this->axis_index);
if (axis_val->uses().size() < 1) {
axis_node->destroy();
}
}

void HandleInitializerNode(std::shared_ptr<Graph> graph, Node* node, Value* axis_val) const {
const std::string initializer_name = axis_val->uniqueName();
for (const auto& initializer : graph->initializers()) {
if (initializer.name() == initializer_name) {
node->i_(kaxis, initializer.int64s().at(0));
node->removeInput(this->axis_index);
// Remove initializer
if (axis_val->uses().size() < 1)
graph->eraseInitializer(initializer_name);
break;
}
}
}

inline Node* EnsureAndReturnNode(Node* node) const {
ONNX_ASSERTM(node->hasAttribute(kaxis), "Axis attribute not created. This may be a bug.");
return node;
}
};

} // namespace version_conversion
} // namespace ONNX_NAMESPACE
6 changes: 5 additions & 1 deletion onnx/version_converter/convert.h
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
#include "onnx/version_converter/BaseConverter.h"
#include "onnx/version_converter/adapters/axes_attribute_to_input.h"
#include "onnx/version_converter/adapters/axes_input_to_attribute.h"
#include "onnx/version_converter/adapters/axis_attribute_to_input.h"
#include "onnx/version_converter/adapters/axis_input_to_attribute.h"
#include "onnx/version_converter/adapters/batch_normalization_13_14.h"
#include "onnx/version_converter/adapters/broadcast_backward_compatibility.h"
#include "onnx/version_converter/adapters/broadcast_forward_compatibility.h"
Expand Down Expand Up @@ -575,9 +577,10 @@ class DefaultVersionConverter : public BaseVersionConverter {
registerAdapter(std::make_unique<CompatibleAdapter>("Size", OpSetID(18), OpSetID(19)));

/******** 19 -> 20 ********/
registerAdapter(std::make_unique<AxisAttributeToInput>("DFT", OpSetID(19), OpSetID(20), 2, 1));
registerAdapter(std::make_unique<CompatibleAdapter>("ConstantOfShape", OpSetID(19), OpSetID(20)));
registerAdapter(std::make_unique<CompatibleAdapter>("IsNaN", OpSetID(19), OpSetID(20)));
registerAdapter(std::make_unique<CompatibleAdapter>("IsInf", OpSetID(19), OpSetID(20)));
registerAdapter(std::make_unique<CompatibleAdapter>("IsNaN", OpSetID(19), OpSetID(20)));
registerAdapter(std::make_unique<CompatibleAdapter>("ReduceMax", OpSetID(19), OpSetID(20)));
registerAdapter(std::make_unique<CompatibleAdapter>("ReduceMin", OpSetID(19), OpSetID(20)));
registerAdapter(std::make_unique<GridSample_19_20>());
Expand All @@ -597,6 +600,7 @@ class DefaultVersionConverter : public BaseVersionConverter {
TensorProto_DataType_FLOAT8E5M2,
TensorProto_DataType_FLOAT8E5M2FNUZ};
registerAdapter(std::make_unique<TypeRestriction>("IsInf", OpSetID(19), OpSetID(20), is_inf_10_unallowed_types));
registerAdapter(std::make_unique<AxisInputToAttribute>("DFT", OpSetID(20), OpSetID(19), 2, -2));
const std::vector<TensorProto_DataType> reduce_min_max_18_unallowed_types = {TensorProto_DataType_BOOL};
registerAdapter(
std::make_unique<TypeRestriction>("ReduceMax", OpSetID(20), OpSetID(19), reduce_min_max_18_unallowed_types));
Expand Down

0 comments on commit e11dacf

Please sign in to comment.