llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[onnx] Fix importer variable names to make `mlir` legal (#2690) 

Some names for `onnx` identifiers are not legal in `mlir-ir`. Sanitize
so that the generated `ir` is legal.
pull/2691/head
Rob Suderman 2023-12-21 17:05:18 -08:00 committed by GitHub
parent ccd469ca0d
commit 85b86b36a2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
python/torch_mlir/extras/onnx_importer.py
View File

@ -38,6 +38,7 @@ from typing import Optional
from dataclasses import dataclass from dataclasses import dataclass
import numpy as np import numpy as np
import re
from ..ir import ( from ..ir import (
ArrayAttr, ArrayAttr,
@ -464,13 +465,18 @@ class ContextCache:
# See TypeProto: sequence_type, map_type, optional_type, sparse_tensor_type. # See TypeProto: sequence_type, map_type, optional_type, sparse_tensor_type.
raise OnnxImportError(f"Unsupported ONNX TypeProto: {tp}") raise OnnxImportError(f"Unsupported ONNX TypeProto: {tp}")
def _sanitize_name(self, name):
if not name.isidentifier():
name = "_" + name
return re.sub("[:/]", "_", name)
def tensor_proto_to_attr(self, tp: onnx.TensorProto) -> Attribute: def tensor_proto_to_attr(self, tp: onnx.TensorProto) -> Attribute:
tensor_type = self.tensor_proto_to_builtin_type(tp) tensor_type = self.tensor_proto_to_builtin_type(tp)
if tp.HasField("raw_data"): if tp.HasField("raw_data"):
# Conveniently, DenseResourceElementsAttr shares the raw data # Conveniently, DenseResourceElementsAttr shares the raw data
# format. We just give it maximum numeric alignment. # format. We just give it maximum numeric alignment.
return DenseResourceElementsAttr.get_from_buffer( return DenseResourceElementsAttr.get_from_buffer(
tp.raw_data, tp.name, tensor_type, alignment=8 tp.raw_data, self._sanitize_name(tp.name), tensor_type, alignment=8
) )
else: else:
# We have to do a data type specific instantiation from proto fields. # We have to do a data type specific instantiation from proto fields.