Merge pull request #53 from tensil-ai/peter/sc-446/support-input-shape-and-onnx-frontend-improvements

petrohi · web-flow · commit e1062f492498 · 2022-05-06T21:27:35.000-07:00
diff --git a/sim/test/src/zynq/tcu/AXIWrapperTCUSpec.scala b/sim/test/src/zynq/tcu/AXIWrapperTCUSpec.scala
@@ -25,7 +25,7 @@ import tensil.{
 }
 import tensil.data.InstructionReader
 
-import tensil.tools.{Compiler, CompilerOptions}
+import tensil.tools.{Compiler, CompilerOptions, CompilerInputShapes}
 import tensil.tools.compiler.MemoryAddressHelper
 import tensil.{InstructionLayout}
 
@@ -100,7 +100,7 @@ class AXIWrapperTCUSpec extends FunUnitSpec {
                 // compiler parameters
                 val options = CompilerOptions(
                   arch = arch,
-                  inputBatchSize = batchSize
+                  inputShapes = CompilerInputShapes.mkWithBatchSize(batchSize),
                 )
 
                 // setup compiler input/output streams
@@ -225,7 +225,7 @@ class AXIWrapperTCUSpec extends FunUnitSpec {
                 // compiler parameters
                 val options = CompilerOptions(
                   arch = arch,
-                  inputBatchSize = batchSize,
+                  inputShapes = CompilerInputShapes.mkWithBatchSize(batchSize),
                   //printProgramFileName = Some(s"sim_resnet20v2_cifar.tasm"),
                 )
 
diff --git a/tools/src/tensil/tools/CompilerInputShapes.scala b/tools/src/tensil/tools/CompilerInputShapes.scala
@@ -0,0 +1,39 @@
+package tensil.tools
+
+object CompilerInputShapes {
+  def parse(shapeStrings: String): CompilerInputShapes = {
+    shapeStrings
+      .split("]\\S*,")
+      .map(_.trim())
+      .map(shapeString => {
+        val Seq(nameString, dimsString) =
+          shapeString.split('[').map(_.trim()).toSeq
+
+        val shape = dimsString
+          .split(']')(0)
+          .split(',')
+          .map(_.trim())
+          .map(dimString =>
+            if (dimString.isEmpty())
+              None
+            else
+              Some(Integer.parseInt(dimString))
+          )
+          .toSeq
+
+        val name =
+          if (nameString.isEmpty())
+            None
+          else
+            Some(nameString)
+
+        (name, shape)
+      })
+      .toMap
+  }
+
+  def mkWithBatchSize(batchSize: Int): CompilerInputShapes =
+    Map(
+      None -> Seq(Some(batchSize))
+    )
+}
diff --git a/tools/src/tensil/tools/CompilerOptions.scala b/tools/src/tensil/tools/CompilerOptions.scala
@@ -12,7 +12,7 @@ case class TracepointCondition(
 
 case class CompilerOptions(
     arch: Architecture,
-    inputBatchSize: Int = 1,
+    inputShapes: CompilerInputShapes = CompilerInputShapes.mkWithBatchSize(1),
     printSummary: Boolean = false,
     printLayersSummary: Boolean = false,
     printSchedulerSummary: Boolean = false,
diff --git a/tools/src/tensil/tools/Main.scala b/tools/src/tensil/tools/Main.scala
@@ -13,7 +13,7 @@ case class Args(
     archFile: File = new File("."),
     modelFile: File = new File("."),
     outputNodes: Seq[String] = Seq("Identity"),
-    inputBatchSize: Int = 1,
+    inputShapes: String = "[1]",
     verbose: Boolean = false,
     summary: Boolean = false,
     layersSummary: Boolean = false,
@@ -41,14 +41,16 @@ object Main extends App {
       .text("Tensil architecture descrition (.tarch) file")
 
     opt[Seq[String]]('o', "output")
-      .valueName("<names>")
+      .valueName("<name>, ...")
       .action((x, c) => c.copy(outputNodes = x))
       .text("Optional list of output nodes, defaults to \"Identity\"")
 
-    opt[Int]('b', "batch")
-      .valueName("<integer>")
-      .action((x, c) => c.copy(inputBatchSize = x))
-      .text("Optional size of input batch, defaults to 1")
+    opt[String]('i', "input-shapes")
+      .valueName("<name> [<dim>, ...], ...")
+      .action((x, c) => c.copy(inputShapes = x))
+      .text(
+        "Optional input shapes, defaults to \"[1]\" (batch size of 1). The shape without <name> is a default for inputs that were not listed by name"
+      )
 
     opt[Boolean]('v', "verbose")
       .valueName("true|false")
@@ -96,7 +98,7 @@ object Main extends App {
 
       val options = CompilerOptions(
         arch = arch,
-        inputBatchSize = args.inputBatchSize,
+        inputShapes = CompilerInputShapes.parse(args.inputShapes),
         printProgress = args.verbose,
         printSummary = args.summary,
         printLayersSummary = args.layersSummary,
diff --git a/tools/src/tensil/tools/compiler/OnnxFrontend.scala b/tools/src/tensil/tools/compiler/OnnxFrontend.scala
@@ -12,7 +12,8 @@ import onnx.onnx.{NodeProto, ModelProto, TensorProto, ValueInfoProto}
 import _root_.tensil.tools.{
   CompilerException,
   TracepointCondition,
-  CompilerOptions
+  CompilerOptions,
+  CompilerInputShapesHelper
 }
 import _root_.tensil.tools.data.{Shape, TensorData}
 import _root_.tensil.tools.util
@@ -579,11 +580,10 @@ class OnnxFrontend(
 
   private def emitInput(context: EmitContext): EmitResult = {
     for ((name, valueInfoProto) <- inputValueInfoProtos) {
-      val shape = Shape(
+      val modelInputShape =
         valueInfoProto.`type`.get.value.tensorType.get.shape.get.dim
-          .map(_.value.dimValue.get.toInt)
-          .toArray
-      )
+          .map(_.value.dimValue.map(_.toInt))
+      val shape = options.inputShapes.deduceInputShape(name, modelInputShape)
 
       val consumers = inputNodeNames(name)
 
@@ -1844,6 +1844,9 @@ class OnnxFrontend(
       scheduler: Scheduler,
       conv2DProto: NodeProto
   ): MemoryObject = {
+    val autoPadAttr = getAttr(conv2DProto, "auto_pad")
+    val autoPad     = autoPadAttr.map(_.s.get.toStringUtf8())
+
     val padsAttr = getAttr(conv2DProto, "pads")
 
     val pads = if (padsAttr.isDefined) {
@@ -1880,15 +1883,6 @@ class OnnxFrontend(
         )
     }
 
-    val (paddingTop, paddingLeft, paddingBottom, paddingRight) =
-      pads.map(_.toInt) match {
-        case Seq(t, l, b, r) => (t, l, b, r)
-        case _ =>
-          throw new CompilerException(
-            s"Unsupported pads [${pads.mkString(", ")}]"
-          )
-      }
-
     context.mm.addPendingConst(
       conv2DProto.input(1),
       getTensorData(tensorProtos(conv2DProto.input(1)))
@@ -1907,6 +1901,45 @@ class OnnxFrontend(
         else None,
       )
 
+    val (paddingTop, paddingLeft, paddingBottom, paddingRight) =
+      pads.map(_.toInt) match {
+        case Seq(t, l, b, r) =>
+          val paddingWidth =
+            (weights.dims.width.toDouble - 1) / 2
+          val paddingHeight =
+            (weights.dims.height.toDouble - 1) / 2
+
+          autoPad match {
+            case Some("SAME_UPPER") =>
+              (
+                Math.floor(paddingHeight).toInt,
+                Math.floor(paddingWidth).toInt,
+                Math.ceil(paddingHeight).toInt,
+                Math.ceil(paddingWidth).toInt
+              )
+
+            case Some("SAME_LOWER") =>
+              (
+                Math.ceil(paddingHeight).toInt,
+                Math.ceil(paddingWidth).toInt,
+                Math.floor(paddingHeight).toInt,
+                Math.floor(paddingWidth).toInt
+              )
+
+            case None | Some("NOTSET") => (t, l, b, r)
+            case Some(v) =>
+              throw new CompilerException(
+                s"Unsupported auto_pad attribute $v"
+              )
+
+          }
+
+        case _ =>
+          throw new CompilerException(
+            s"Unsupported pads [${pads.mkString(", ")}]"
+          )
+      }
+
     val inputVars =
       context.mm.consumeObject(conv2DProto.input(0), Seq(conv2DProto.name.get))
 
@@ -2471,7 +2504,15 @@ class OnnxFrontend(
     val input1Name =
       if (addProto.input(0) == input0Temp.name) addProto.input(1)
       else addProto.input(0)
-    val input1Vars =
+
+    val input1Vars = if (tensorProtos.isDefinedAt(input1Name)) {
+      context.mm.addPendingConst(
+        input1Name,
+        getTensorData(tensorProtos(input1Name))
+      )
+
+      context.mm.getOrEmitConstObject(input1Name)
+    } else
       context.mm.consumeObject(input1Name, Seq(addProto.name.get))
 
     scheduler.emitAdd(
diff --git a/tools/src/tensil/tools/compiler/TfFrontend.scala b/tools/src/tensil/tools/compiler/TfFrontend.scala
@@ -14,7 +14,8 @@ import org.tensorflow.framework.types.DataType
 import _root_.tensil.tools.{
   CompilerException,
   TracepointCondition,
-  CompilerOptions
+  CompilerOptions,
+  CompilerInputShapesHelper
 }
 import _root_.tensil.tools.data.{Shape, TensorData}
 import _root_.tensil.tools.util
@@ -575,16 +576,12 @@ class TfFrontend(
       context: EmitContext,
       placeholderDef: NodeDef
   ): EmitResult = {
-    val shape = util.getShape(placeholderDef)
-
+    val modelInputShape = util
+      .getShape(placeholderDef)
+      .map(dim => if (dim >= 0) Some(dim) else None)
+      .toSeq
     val placeholderShape =
-      if (shape(0) == -1)
-        Shape(
-          options.inputBatchSize +: shape
-            .takeRight(shape.size - 1)
-            .toArray
-        )
-      else shape
+      options.inputShapes.deduceInputShape(placeholderDef.name, modelInputShape)
 
     val placeholderDims = VarsDimensions(placeholderShape)
 
diff --git a/tools/src/tensil/tools/package.scala b/tools/src/tensil/tools/package.scala
@@ -4,8 +4,58 @@
 package tensil
 
 import tensil.tools.compiler.{MemoryAddress, MemoryObject}
+import tensil.tools.data.Shape
 
 package object tools {
-  type TracepointsMap = Map[MemoryAddress, List[MemoryObject]]
-  type CompilerSourceType = String
+  type TracepointsMap      = Map[MemoryAddress, List[MemoryObject]]
+  type CompilerSourceType  = String
+  type CompilerInputShape  = Seq[Option[Int]]
+  type CompilerInputShapes = Map[Option[String], CompilerInputShape]
+
+  implicit class CompilerInputShapeHelper(val inputShape: CompilerInputShape) {
+    override def toString() =
+      s"[${inputShape.map(v => if (v.isDefined) v.get.toString else "?").mkString(", ")}]"
+  }
+
+  implicit class CompilerInputShapesHelper(
+      val inputShapes: CompilerInputShapes
+  ) {
+    def batchSize = inputShapes.head._2(0).get
+
+    def deduceInputShape(
+        name: String,
+        modelInputShape: CompilerInputShape
+    ): Shape = {
+      Shape(
+        modelInputShape.zipWithIndex
+          .map({
+            case (modelDim, i) =>
+              val optionsInputShape = inputShapes
+                .getOrElse(Some(name), inputShapes(None))
+
+              if (
+                optionsInputShape
+                  .isDefinedAt(i) && optionsInputShape(i).isDefined
+              ) {
+                val optionsDim = optionsInputShape(i).get
+
+                if (modelDim.isDefined && modelDim.get != optionsDim)
+                  throw new CompilerException(
+                    s"Specified input shape for $name ${CompilerInputShapeHelper(optionsInputShape)} is incompatible with model shape ${CompilerInputShapeHelper(modelInputShape)}"
+                  )
+
+                optionsDim
+              } else {
+                if (modelDim.isDefined)
+                  modelDim.get
+                else
+                  throw new CompilerException(
+                    s"Specified input shape for $name ${CompilerInputShapeHelper(optionsInputShape)} has unspecified dimensions in model shape ${CompilerInputShapeHelper(modelInputShape)}"
+                  )
+              }
+          })
+          .toArray
+      )
+    }
+  }
 }
diff --git a/tools/test/src/tools/CompilerSpec.scala b/tools/test/src/tools/CompilerSpec.scala
