Skip to content

Commit 9bbe1d6

Browse files
Thiago Crepaldipytorchmergebot
Thiago Crepaldi
authored and
pytorchmergebot
committed
Fix ONNX ATen fallback for non-caffe2 engines
This PR introduces 3 BC changes: First, this PR propagates `BUILD_CAFFE2` flag to `libtorch` and `libtorch_python`, which is necessary for non-caffe2 ONNX runtimes when using `ONNX_ATEN_FALLBACK` operator export type. Second, as a complement of pytorch#68490, this PR refactors Caffe2's Aten ops symbolics to consider not only the `operator_export_type` (aka `ONNX_ATEN_FALLBACK`) to emit Caffe2 Aten ops, but also whether `BUILD_CAFFE2` (which is called `torch.onnx._CAFFE2_ATEN_FALLBACK` in python binding) is set. Lastly, it renames `onnx::ATen` to `aten::ATen` for ONNX spec consistency in a BC fashion. ONNX doesn't have `ATen` op on its spec, but PyTorch ONNX converter emits them. Non-Caffe2 backend engines would be mislead by such operator's name/domain. A non-ideal workaround would be to have Aten ops handled based on its name and ignore the (non-complaint) domain. Moreover, users could incorrectly file bugs to either ONNX or ONNX Runtime when they inspect the model and notice the presence of an unspecified ONNX operator. Pull Request resolved: pytorch#73954 Approved by: https://github.com/BowenBao, https://github.com/malfet, https://github.com/garymm, https://github.com/jiafatom
1 parent b142a22 commit 9bbe1d6

26 files changed

+146
-112
lines changed

‎aten/src/ATen/core/interned_strings.h

+3-2
Original file line numberDiff line numberDiff line change
@@ -247,7 +247,7 @@ namespace c10 {
247247
_(onnx, Less) \
248248
_(onnx, LessOrEqual) \
249249
_(onnx, Not) \
250-
_(onnx, ATen) \
250+
_(aten, ATen) \
251251
_(onnx, Split) \
252252
_(onnx, ConstantOfShape) \
253253
_(onnx, Cast) \
@@ -316,7 +316,8 @@ namespace c10 {
316316
_(attr, new_axis) \
317317
_(attr, warn_id) \
318318
_(attr, allowzero) \
319-
_(attr, seen_none)
319+
_(attr, seen_none) \
320+
_(attr, overload_name)
320321

321322
enum class _keys : unique_t {
322323
#define DEFINE_KEY(ns, s) ns##_##s,

‎binaries/bench_gen/bench_gen.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ def main(args):
5959

6060
if __name__ == "__main__":
6161
parser = argparse.ArgumentParser(
62-
description="Utilitity to generate Caffe2 benchmark models.")
62+
description="Utility to generate Caffe2 benchmark models.")
6363
parser.add_argument("operator", help="Caffe2 operator to benchmark.")
6464
parser.add_argument("-b", "--blob",
6565
help="Instantiate a blob --blob name=dim1,dim2,dim3",

‎caffe2/CMakeLists.txt

+2
Original file line numberDiff line numberDiff line change
@@ -1943,6 +1943,8 @@ if(BUILD_PYTHON)
19431943
# ---[ Python.
19441944
if(BUILD_CAFFE2)
19451945
add_library(caffe2_pybind11_state MODULE ${Caffe2_CPU_PYTHON_SRCS})
1946+
target_compile_definitions(torch PRIVATE BUILD_CAFFE2)
1947+
target_compile_definitions(torch_python PRIVATE BUILD_CAFFE2)
19461948
if(USE_NUMPY)
19471949
target_compile_options(caffe2_pybind11_state PRIVATE "-DUSE_NUMPY")
19481950
target_link_libraries(caffe2_pybind11_state PRIVATE numpy::numpy)

‎caffe2/contrib/aten/README.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ class Add(torch.autograd.Function):
7272
7373
@staticmethod
7474
def symbolic(g, a, b):
75-
return g.op("ATen", a, b, operator_s = "add")
75+
return g.at("add", a, b)
7676
7777
@staticmethod
7878
def forward(ctx, a, b):

‎caffe2/contrib/aten/aten_op_template.h

+1-1
Original file line numberDiff line numberDiff line change
@@ -179,7 +179,7 @@ class ATenOp : public Operator<Context> {
179179
std::vector<std::string> attrs;
180180
for (const auto i : c10::irange(operator_def.arg_size())) {
181181
auto & attr = operator_def.arg(i);
182-
if(attr.name() == "operator" || attr.name() == "type" || attr.name() == "overload_name" ) {
182+
if (attr.name() == "operator" || attr.name() == "type" || attr.name() == "overload_name") {
183183
continue;
184184
}
185185
attrs.push_back(attr.name());

‎caffe2/contrib/aten/aten_test.py

+1-6
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,4 @@
1-
2-
3-
4-
5-
6-
from caffe2.python import core, dyndep
1+
from caffe2.python import core
72
from hypothesis import given
83

94
import caffe2.python.hypothesis_test_util as hu

‎caffe2/contrib/aten/docs/sample.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,8 @@ def forward(self, x, y):
3838
# graph(%input : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu),
3939
# %y : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu)):
4040
# %2 : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu) = onnx::Relu(%input)
41-
# %3 : Tensor = onnx::ATen[operator="mul"](%2, %2)
42-
# %4 : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu) = onnx::ATen[operator="add"](%3, %y)
41+
# %3 : Tensor = aten::ATen[operator="mul"](%2, %2)
42+
# %4 : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu) = aten::ATen[operator="add"](%3, %y)
4343
# return (%4)
4444

4545
graph = onnx.load(f.name)

‎caffe2/python/benchmark_generator.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,7 @@ def make_blob_on_context(blob_name, blob_data, context):
106106

107107
if __name__ == "__main__":
108108
parser = argparse.ArgumentParser(
109-
description="Utilitity to generate Caffe2 benchmark models.")
109+
description="Utility to generate Caffe2 benchmark models.")
110110
parser.add_argument("operator", help="Caffe2 operator to benchmark.")
111111
parser.add_argument("-b", "--blob",
112112
help="Instantiate a blob --blob name=dim1,dim2,dim3",

‎test/expect/TestPytorchExportModes.test_aten_fallback.expect

+1-1
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ ModelProto {
1111
nodes: [
1212
Node {type: "Add", inputs: [0,1], outputs: [2], attributes: []},
1313
Node {type: "Constant", inputs: [], outputs: [3], attributes: [{ name: 'value', type: tensor, value:TensorProto shape: []}]},
14-
Node {type: "ATen", inputs: [2,3], outputs: [4,5], attributes: [{ name: 'operator', type: string, value: 'qr'}, { name: 'overload_name', type: string, value: ''}]}
14+
Node {type: "ATen", domain: "org.pytorch.aten", inputs: [2,3], outputs: [4,5], attributes: [{ name: 'operator', type: string, value: 'qr'}, { name: 'overload_name', type: string, value: ''}]}
1515
]
1616
}
1717
opset_import: [OperatorSetIdProto { domain: }OperatorSetIdProto { domain: org.pytorch.aten}],

‎test/expect/TestPytorchExportModes.test_onnx_aten.expect

+1-1
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ ModelProto {
99
outputs: [{name: "2", type:Tensor dims: 3 4}]
1010
initializers: []
1111
nodes: [
12-
Node {type: "ATen", inputs: [0,1], outputs: [2], attributes: [{ name: 'operator', type: string, value: 'fmod'}, { name: 'overload_name', type: string, value: ''}]}
12+
Node {type: "ATen", domain: "org.pytorch.aten", inputs: [0,1], outputs: [2], attributes: [{ name: 'operator', type: string, value: 'fmod'}, { name: 'overload_name', type: string, value: ''}]}
1313
]
1414
}
1515
opset_import: [OperatorSetIdProto { domain: }OperatorSetIdProto { domain: org.pytorch.aten}],

‎test/expect/TestScript.test_listconstruct_erasure.expect

+1-1
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ ModelProto {
1313
Node {type: "Less", inputs: [0,1], outputs: [2], attributes: []},
1414
Node {type: "Cast", inputs: [2], outputs: [3], attributes: [{ name: 'to', type: int, value: 2}]},
1515
Node {type: "Cast", inputs: [3], outputs: [4], attributes: [{ name: 'to', type: int, value: 9}]},
16-
Node {type: "ATen", inputs: [0,4], outputs: [5], attributes: [{ name: 'operator', type: string, value: 'index'}, { name: 'overload_name', type: string, value: ''}]}
16+
Node {type: "ATen", domain: "org.pytorch.aten", inputs: [0,4], outputs: [5], attributes: [{ name: 'operator', type: string, value: 'index'}, { name: 'overload_name', type: string, value: ''}]}
1717
]
1818
}
1919
opset_import: [OperatorSetIdProto { domain: }OperatorSetIdProto { domain: org.pytorch.aten}],

‎test/onnx/expect/TestOperators.test_at_op.expect

+5
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ graph {
1818
s: ""
1919
type: STRING
2020
}
21+
domain: "org.pytorch.aten"
2122
}
2223
name: "torch_jit"
2324
input {
@@ -56,3 +57,7 @@ graph {
5657
opset_import {
5758
version: 13
5859
}
60+
opset_import {
61+
domain: "org.pytorch.aten"
62+
version: 1
63+
}

‎test/onnx/expect/TestOperators.test_aten_embedding_2.expect

+11-22
Original file line numberDiff line numberDiff line change
@@ -6,19 +6,24 @@ graph {
66
input: "emb.weight"
77
input: "input_1"
88
output: "onnx::Add_3"
9-
name: "ATenOp_0"
10-
op_type: "ATenOp"
9+
name: "ATen_0"
10+
op_type: "ATen"
1111
attribute {
1212
name: "custom_attributes_json"
1313
s: "{\"padding_idx\":-1,\"scale_grad_by_freq\":false,\"sparse\":false}"
1414
type: STRING
1515
}
1616
attribute {
17-
name: "name"
18-
s: "aten::embedding"
17+
name: "operator"
18+
s: "embedding"
1919
type: STRING
2020
}
21-
domain: "com.microsoft"
21+
attribute {
22+
name: "overload_name"
23+
s: ""
24+
type: STRING
25+
}
26+
domain: "org.pytorch.aten"
2227
}
2328
node {
2429
input: "onnx::Add_3"
@@ -145,27 +150,11 @@ graph {
145150
}
146151
}
147152
}
148-
value_info {
149-
name: "onnx::Add_3"
150-
type {
151-
tensor_type {
152-
elem_type: 1
153-
shape {
154-
dim {
155-
dim_param: "ATenOponnx::Add_3_dim_0"
156-
}
157-
dim {
158-
dim_param: "ATenOponnx::Add_3_dim_1"
159-
}
160-
}
161-
}
162-
}
163-
}
164153
}
165154
opset_import {
166155
version: 12
167156
}
168157
opset_import {
169-
domain: "com.microsoft"
158+
domain: "org.pytorch.aten"
170159
version: 1
171160
}

0 commit comments

Comments
 (0)