Add inference for Identity operation. (#400)

src/compiler/dialect/onnx/gen_doc.py

@@ -266,7 +266,8 @@ def gen_schema(schema) :
     ShapeInferenceList=['Exp', 'Tanh', 'Sinh', 'Cosh', 'Sigmoid', 'Relu',
                         'Add', 'Mul', 'Div', 'Sub', 'And', 'Or', 'Xor',
                         'Sum', 'Max', 'Min', 'MatMul', 'Gemm', 'LeakyRelu',
-                        'Elu', 'Selu', 'HardSigmoid', 'Reshape', 'Reciprocal']
+                        'Elu', 'Selu', 'HardSigmoid', 'Reshape', 'Reciprocal',
+                        'Identity']
     CanonicalList=['Add', 'Identity']
     line_indent = '  '
 

src/compiler/dialect/onnx/onnx_ops.cpp

@@ -8,8 +8,6 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "llvm/ADT/SetVector.h"
-#include "llvm/ADT/SmallBitVector.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
@@ -17,6 +15,8 @@
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallBitVector.h"
 
 #include "onnx_ops.hpp"
 
@@ -202,6 +202,14 @@ void ONNXMinOp::inferShapes() {
   getResult()->setType(getOperand(0)->getType());
 }
 
+//===----------------------------------------------------------------------===//
+// Identity
+/// Infer the output shape of the ONNXIdentityOp. This method is required by the
+/// shape inference interface.
+void ONNXIdentityOp::inferShapes() {
+  getResult()->setType(getOperand()->getType());
+}
+
 //===----------------------------------------------------------------------===//
 
 // MatMul

src/compiler/dialect/onnx/onnxop.inc

@@ -1026,7 +1026,7 @@ def ONNXHardmaxOp:ONNX_Op<"Hardmax",
 }
 
 def ONNXIdentityOp:ONNX_Op<"Identity", 
-    [NoSideEffect]> {
+    [NoSideEffect, DeclareOpInterfaceMethods<ShapeInferenceOpInterface>]> {
   let hasCanonicalizer = 1;
   let summary = "ONNX Identity operation";
   let description = [{

src/compiler/pass/shape_inference_pass.cpp

@@ -9,9 +9,9 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Pass/Pass.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/raw_ostream.h"
-#include "mlir/Pass/Pass.h"
 
 #include "shape_inference_interface.hpp"
 #include "src/compiler/dialect/onnx/onnx_ops.hpp"
@@ -58,8 +58,7 @@ class ShapeInferencePass : public mlir::FunctionPass<ShapeInferencePass> {
       if (auto shape_op = dyn_cast<ShapeInference>(op)) {
         shape_op.inferShapes();
       } else {
-        op->emitError(
-            "unable to infer shape of operation without shape "
+        op->emitError("unable to infer shape of operation without shape "
                       "inference interface");
         return signalPassFailure();
       }
@@ -74,7 +73,8 @@ class ShapeInferencePass : public mlir::FunctionPass<ShapeInferencePass> {
 
     if (auto terminator_op = f.getBody().back().getTerminator()) {
       auto results = terminator_op->getOperandTypes();
-      f.setType(FunctionType::get(f.getType().getInputs(),
+      f.setType(FunctionType::get(
+          f.getType().getInputs(),
           std::vector<Type>(results.begin(), results.end()), f.getContext()));
     }
   }
@@ -109,13 +109,15 @@ class ShapeInferencePass : public mlir::FunctionPass<ShapeInferencePass> {
         op->getName().getStringRef() != "onnx.Sum" &&
         op->getName().getStringRef() != "onnx.Max" &&
         op->getName().getStringRef() != "onnx.Min" &&
+        op->getName().getStringRef() != "onnx.Identity" &&
         op->getName().getStringRef() != "onnx.MatMul" &&
         op->getName().getStringRef() != "onnx.Gemm" &&
         op->getName().getStringRef() != "onnx.FullGemm" &&
         op->getName().getStringRef() != "onnx.Reshape")
       return false;
-    return llvm::any_of(op->getResultTypes(),
-        [](Type result_type) { return !result_type.isa<RankedTensorType>(); });
+    return llvm::any_of(op->getResultTypes(), [](Type result_type) {
+      return !result_type.isa<RankedTensorType>();
+    });
   }
 };
 } // end anonymous namespace
@@ -127,5 +129,5 @@ std::unique_ptr<mlir::Pass> mlir::createShapeInferencePass() {
   return std::make_unique<ShapeInferencePass>();
 }
 
-static PassRegistration<ShapeInferencePass> pass(
-    "shape-inference", "Shape inference for frontend dialects.");
+static PassRegistration<ShapeInferencePass>
+    pass("shape-inference", "Shape inference for frontend dialects.");