Add peephole to aten::device(str, int)

torch/_dynamo/backends/onnxrt.py

@@ -72,7 +72,6 @@ def onnxrt(gm, example_inputs, *, filename=None, provider=None):
         filename,
         input_names=input_names,
         output_names=output_names,
-        opset_version=17,
     )
     del example_inputs, example_outputs
 

torch/csrc/jit/passes/peephole.cpp

@@ -269,6 +269,24 @@ struct PeepholeOptimizeImpl {
           node->output()->replaceAllUsesWith(output);
           changed = true;
         }
+      } else if (
+          node->matches("aten::device(str type, int index) -> Device") &&
+          shape_peepholes_) {
+        auto string_type = node->inputs().at(0)->type()->expect<StringType>();
+        if (string_type) {
+          WithInsertPoint guard(node);
+          std::string type_str = node->inputs().at(0)->node()->s(attr::value);
+          auto index = toIValue(node->inputs().at(1))->toInt();
+          auto device = c10::Device(type_str + ":" + std::to_string(index));
+          auto output = node->owningGraph()->insertConstant(device);
+          GRAPH_UPDATE(
+              "Replacing ",
+              getHeader(node),
+              " with a device constant ",
+              output->debugName());
+          node->output()->replaceAllUsesWith(output);
+          changed = true;
+        }
       } else if (
           node->matches("aten::dim(Tensor self) -> int") && shape_peepholes_) {
         auto ptt = node->input()->type()->expect<TensorType>();

torch/onnx/symbolic_opset17.py

@@ -65,13 +65,6 @@ def _compute_edge_sizes(n_fft, window_size):
     return left, right
 
 
-# @_onnx_symbolic("aten::device")
-# # @symbolic_helper.parse_args("s", "i")
-# def device(g: jit_utils.GraphContext, type, index):
-#     # device as a noop
-#     return type
-
-
 @_onnx_symbolic("aten::stft")
 @symbolic_helper.parse_args("v", "i", "i", "i", "v", "b", "b", "b")
 @_beartype.beartype