neural_net_dsl_v3.nim

import macros, tables, hashes

import ../src/arraymancer

proc flatten*(s: openarray[int]): int {.inline.}=

  assert s.len != 0

  result = 1

  for val in s:
    result *= val

proc hash(x: NimNode): Hash =
  assert x.kind == nnkIdent
  result = hash($x)

type
  LayerKind = enum
    lkInput, lkConv2D, lkLinear, lkMaxPool2D

  LayerTopology = object
    ## Describe a layer topology
    in_shape, out_shape: NimNode # Input and output shape
    case kind: LayerKind
    of lkConv2D:
      c2d_kernel_shape: NimNode
      c2d_padding, c2d_strides: NimNode
    of lkMaxPool2D:
      m2d_kernel, m2d_padding, m2d_strides: NimNode
    else:
      discard

  TrainableLayer[TT] = object of RootObj
    weight: Variable[TT]
    bias: Variable[TT]

  Conv2DLayer[T] = object of TrainableLayer[T]
  LinearLayer[T] = object of TrainableLayer[T]

  #################################################

  TopoTable = Table[NimNode, LayerTopology]

  NetworkSections = tuple[layers, forward: NimNode]

  ModelField = tuple[field_name: NimNode, field_type: NimNode, init_call: NimNode]
    ## Field name, type and initialization proc

  Neuromancer = ref object
    ## Pathfinder that will go through the network topology
    ## and create the corresponding type, initialization and forward proc

    # Running state
    topoTable: TopoTable          # Retrieve the layer properties.
                                  # Maps a NimNode of nnkIdent to the corresponding input and output shape
    subtype: NimNode

    # Outputs
    context: NimNode                # copy-paste of the context
    type_section: NimNode           # The type section to declare the new Model type
    trainparams: seq[ModelField]    # stores the weight/bias and initialization proc
    init_proc: NimNode              # proc init(T: typedesc[Model]): Model
    forward_proc: NimNode           # proc forward(self: Model, x, y: Variable)
    forward_templates: seq[NimNode] # This adds templates synctactic sugar, so that the field name is substituted
                                    # template cv1(x: Variable[T]): Variable[T] =
                                    #   x.conv2d(self.cv1.weight, self.cv1.bias)
    forward_asserts: NimNode

proc out_shape_conv2d(in_shape: array[3, int], kernel: array[4, int], padding, strides: tuple[h, w: int]): array[3, int] {.noInit.}=
  ## Each dimension of the (nbDims-2)-D images of the output tensor is computed as followed:
  ##   outputDim = 1 + ( inputDim + 2*pad - (((filterDim-1)*upscaleA)+1) )/ convolutionStride;

  ## Input and result are of shape [C, H, W]
  ## Kernel of shape [C_out, C_in, h, w]

  # Reminder we don't consider N (bath size) in the topology
  template kH: int = kernel[2]
  template kW: int = kernel[3]
  template pH: int = padding.h
  template pW: int = padding.w
  template sH: int = strides.h
  template sW: int = strides.w

  template iH: int = in_shape[1]
  template iW: int = in_shape[2]
  template dH: int = 1 # dilation # TODO
  template dW: int = 1 # dilation

  result[0] = kernel[0]                                    # C
  result[1] = 1 + (iH + 2*pH - (((kH-1) * dH) + 1) div sH) # H
  result[2] = 1 + (iW + 2*pW - (((kW-1) * dW) + 1) div sW) # W

proc out_shape_maxpool2d(in_shape: array[3, int], kernel, padding, strides: tuple[h, w: int]): array[3, int] {.noInit.}=

  # Reminder we don't consider N (bath size) in the topology
  template C: int = in_shape[0]
  template H: int = in_shape[1]
  template W: int = in_shape[2]

  template kH: int = kernel.h
  template kW: int = kernel.w
  template pH: int = padding.h
  template pW: int = padding.w
  template sH: int = strides.h
  template sW: int = strides.w

  result[0] = C
  result[1] = (H + (2 * pH) - kH) div sH + 1
  result[2] = (W + (2 * pW) - kW) div sW + 1

proc splitSections(config: NimNode): NetworkSections =
  template unknown =
    error:
      lineInfo(section) &
        ": unknown neural network configuration section \"" &
        $section[0] & "\""

  for section in config:
    if section.kind == nnkCall:
      if eqIdent(section[0], "layers"):
        result.layers = section[1]
      else:
        unknown()
    elif section.kind == nnkCommand:
      if eqIdent(section[0], "forward"):
        # For forward we copy everything.
        # We have to deal with forward with multiple inputs like "forward x, y, z:"
        # and we will do that later.
        result.forward = section
      else:
        unknown()
    else:
        unknown()

proc isValidLayerSection(section: NimNode): bool =

  # Expected AST - cv1: Conv2D(20, 5, 5)
  # Call
  #   Ident "cv1"
  #   StmtList
  #     Call
  #       Ident "Conv2D"
  #       IntLit 20
  #       IntLit 5
  #       IntLit 5
  (section.kind == nnkCall) and
    (section[0].kind == nnkIdent) and
    (section[1].kind == nnkStmtlist) and
    (section[1].len == 1) and
    (section[1][0].kind == nnkCall)

template unknown(section: Nimnode) =
  error:
    lineInfo(section) &
      ": unknown neural network configuration section \"" &
      $section[0] & "\""

template incorrect(section: Nimnode) =
  error:
    lineInfo(section) &
      ": incorrect neural network configuration section \"" &
      $section[0] & "\""

proc topoFromInput(self: var TopoTable, ident: NimNode, desc: NimNode) =

  # Initializes the ident --> (input, output) topology table with the input shapes

  # Call
  #   Ident "Input"
  #   Bracket
  #     IntLit 1
  #     IntLit 28
  #     IntLit 28

  if desc.len != 2:
    incorrect(desc) ## Placeholder to specify padding stride in the future

  self.add ident, LayerTopology(kind: lkInput,
                                in_shape: desc[1],
                                out_shape: desc[1])

template letsGoDeeper =
  var rTree = node.kind.newTree()
  for child in node:
    rTree.add inspect(child)
  return rTree

proc replaceInputNodes(self: TopoTable, in_shape: NimNode): NimNode =
  # Args:
  #   - The topology table
  #   - the input shape
  # Returns:
  #   - An AST input shape with "x.out_shape" replaced by the actual x.out_shape
  #     taken from the topology table

  proc inspect(node: NimNode): NimNode =
    case node.kind:
    of nnkDotExpr:
      if eqIdent(node[1], "out_shape"):
        return self[node[0]].out_shape
      else:
        letsGoDeeper()
    of {nnkIdent, nnkSym, nnkEmpty}:
      return node
    of nnkLiterals:
      return node
    else:
      letsGoDeeper()
  result = inspect(in_shape)

proc replaceSymsByIdents*(ast: NimNode): NimNode =
  proc inspect(node: NimNode): NimNode =
    case node.kind:
    of {nnkIdent, nnkSym}:
      return ident($node)
    of nnkEmpty:
      return node
    of nnkLiterals:
      return node
    else:
      letsGoDeeper()
  result = inspect(ast)

proc topoFromConv2D(self: var TopoTable, ident: NimNode, desc: NimNode) =

  # Call
  #   Ident "Conv2D"
  #   # Kernel (C_out, kH, kW)
  #   IntLit 20
  #   IntLit 5
  #   IntLit 5
  #   # Kernel strides & padding

  var padding, strides: NimNode

  if desc.len > 5:
    incorrect(desc) ## Placeholder to specify padding stride in the future
  else:
    padding = quote do: (0, 0)
    strides = quote do: (1, 1)

  var in_shape = self.replaceInputNodes(desc[1])
  in_shape = quote do: `in_shape`

  let
    c_out = desc[2]
    kH = desc[3]
    kW = desc[4]

  let kernel = quote do:
    # C_out, C_in, kH, kW
    [`c_out`, `in_shape`[0], `kH`, `kW`]

  let out_shape = quote do:
    out_shape_conv2d(`in_shape`, `kernel`, `padding`, `strides`)

  self.add ident, LayerTopology(kind: lkConv2D,
                                in_shape: in_shape,
                                out_shape: out_shape,
                                c2d_kernel_shape: kernel,
                                c2d_padding: padding,
                                c2d_strides: strides)

proc topoFromMaxPool2D(self: var TopoTable, ident: NimNode, desc: NimNode) =

  # Call
  #   Ident "MaxPool2D"
  #   Par
  #     IntLit 2
  #     IntLit 2
  #   Par
  #     IntLit 0
  #     IntLit 0
  #   Par
  #     IntLit 2
  #     IntLit 2

  if desc.len != 5:
    incorrect(desc) ## Placeholder to specify padding stride in the future

  var in_shape = self.replaceInputNodes(desc[1])
  in_shape = quote do: `in_shape`

  let
    kernel = desc[2]
    padding = desc[3]
    strides = desc[4]

  let out_shape = quote do:
    out_shape_maxpool2d(`in_shape`, `kernel`, `padding`, `strides`)

  self.add ident, LayerTopology(kind: lkMaxPool2D,
                                in_shape: in_shape,
                                out_shape: out_shape,
                                m2d_kernel: kernel,
                                m2d_padding: padding,
                                m2d_strides: strides)

proc topoFromLinear(self: var TopoTable, ident: NimNode, desc: NimNode) =

  # Call
  #   Ident "Linear"
  #   IntLit 10

  if desc.len != 3:
    incorrect(desc) ## Placeholder to specify padding stride in the future

  var in_shape = self.replaceInputNodes(desc[1])
  in_shape = quote do: `in_shape`

  self.add ident, LayerTopology(kind: lkLinear,
                                in_shape: in_shape,
                                out_shape: desc[2])

proc topoFromLayer(self: var TopoTable, ident: NimNode, desc: NimNode) =

  if eqIdent(desc[0], "Conv2D"):
    self.topoFromConv2D(ident, desc)
  elif eqIdent(desc[0], "MaxPool2D"):
    self.topoFromMaxPool2D(ident, desc)
  elif eqIdent(desc[0], "Linear"):
    self.topoFromLinear(ident, desc)
  elif eqIdent(desc[0], "Input"):
    self.topoFromInput(ident, desc)
  else:
    unknown(desc)

proc topoFromLayers(self: var TopoTable, layers: NimNode) =

  ## Add all layers and their known parameters to the table
  #

  for section in layers:
    if section.isValidLayerSection:
      assert section[0] notin self
      self.topoFromLayer(
        ident($section[0]),
        section[1][0]
        )
    else:
      unknown(section)

proc trainParamsConv2D(self: Neuromancer, field_name: NimNode, topo: LayerTopology) =

  # 1. Create the object field
  var convConfig: ModelField

  convConfig.field_name = field_name
  convConfig.field_type = nnkBracketExpr.newTree(
      ident("Conv2DLayer"), self.subtype
    )

  # 2. Configure weight and bias
  let
    topo = self.topoTable.getOrDefault(field_name)
    sst = self.subtype[1] # we need to get the subsubtype float32/float64

    kshape = topo.c2d_kernel_shape

    w_shape = kshape
    b_shape = quote do: [`kshape`[0], 1, 1]

    w = quote do: randomTensor(`w_shape`, `sst`(1)) .- `sst`(0.5)
    b = quote do: randomTensor(`b_shape`, `sst`(1)) .- `sst`(0.5)

  convConfig.init_call = newStmtList()

  let ctx = self.context

  convConfig.init_call.add quote do:
    result.`field_name`.weight = `ctx`.variable(
      `w`, requires_grad = true # TODO allow freezing
    )

    result.`field_name`.bias = `ctx`.variable(
      `b`, requires_grad = true # TODO allow freezing
    )

  self.trainparams.add convConfig

proc trainParamsLinear(self: Neuromancer, field_name: NimNode, topo: LayerTopology) =

  # 1. Create the object field
  var linearConfig: ModelField

  linearConfig.field_name = field_name
  linearConfig.field_type = nnkBracketExpr.newTree(
      ident("LinearLayer"), self.subtype
    )

  # 2. Configure weight and bias
  let
    topo = self.topoTable.getOrDefault(field_name)
    sst = self.subtype[1] # we need to get the subsubtype float32/float64

    in_shape = topo.in_shape
    out_shape = topo.out_shape

    w_shape = quote do: [`out_shape`, `in_shape`]
    b_shape = quote do: [1, `out_shape`]

    w = quote do: randomTensor(`w_shape`, `sst`(1)) .- `sst`(0.5)
    b = quote do: randomTensor(`b_shape`, `sst`(1)) .- `sst`(0.5)

  linearConfig.init_call = newStmtList()

  let ctx = self.context

  linearConfig.init_call.add quote do:
    result.`field_name`.weight = `ctx`.variable(
      `w`, requires_grad = true # TODO allow freezing
    )

    result.`field_name`.bias = `ctx`.variable(
      `b`, requires_grad = true # TODO allow freezing
    )

  self.trainparams.add linearConfig

proc genModelFieldInit(self: Neuromancer) =

  self.trainparams = @[]

  for k, v in pairs(self.topoTable):
    case v.kind:
    of lkConv2D: self.trainParamsConv2D(k, v)
    of lkLinear: self.trainParamsLinear(k, v)
    else:
      discard

proc shortcutConv2D(self: Neuromancer, field_name: NimNode, topo: LayerTopology) =

  # TODO: Add padding and strides

  let shortcut = quote do:
    template `field_name`(x: Variable): Variable =
      x.conv2d(self.`field_name`.weight, self.`field_name`.bias)

  self.forward_templates.add shortcut

proc shortcutMaxPool2D(self: Neuromancer, field_name: NimNode, topo: LayerTopology) =

  let
    topo = self.topoTable.getOrDefault(field_name)
    kernel = topo.m2d_kernel
    padding = topo.m2d_padding
    strides = topo.m2d_strides

  let shortcut = quote do:
    template `field_name`(x: Variable): Variable =
      x.maxpool2D(`kernel`, `padding`, `strides`)

  self.forward_templates.add shortcut

proc shortcutLinear(self: Neuromancer, field_name: NimNode, topo: LayerTopology) =

  let shortcut = quote do:
    template `field_name`(x: Variable): Variable =
      x.conv2d(self.`field_name`.weight, self.`field_name`.bias)

  self.forward_templates.add shortcut

proc genTemplateShortcuts(self: Neuromancer) =

  self.forward_templates = @[]

  for k, v in pairs(self.topoTable):
    case v.kind:
    of lkConv2D: self.shortcutConv2D(k, v)
    of lkLinear: self.shortcutLinear(k, v)
    of lkMaxPool2D: self.shortcutMaxPool2D(k, v)
    else:
      discard

proc genModelType(self: Neuromancer, model_name: string) =

  var records = nnkRecList.newTree

  for record in self.trainparams:
    let (field_name, field_type, _) = record
    records.add nnkIdentDefs.newTree(
      newIdentNode($field_name),
      field_type,
      newEmptyNode()
    )

  self.type_section = nnkStmtList.newTree(
    nnkTypeSection.newTree(
      nnkTypeDef.newTree(
        newIdentNode(model_name),
        newEmptyNode(), # Generic params get here
        nnkObjectTy.newTree(
          newEmptyNode(),
          newEmptyNode(),
          records
        )
      )
    )
  )

proc genInitProc(self: Neuromancer, model_name: string) =

  self.init_proc = newStmtList()

  let
    subtype = self.subtype
    modelType = newIdentNode(model_name)
    procBody = newStmtList()

  for record in self.trainparams:
    let (_, _, initStmt) = record
    procBody.add initStmt

  self.init_proc.add quote do:
    proc init(ctx: Context[`subtype`], model_type: typedesc[`modelType`]): `modelType` =
      `procBody`

proc genForwardProc(self: Neuromancer, model_name: string, forward: NimNode) =

  forward.expectKind(nnkCommand)
  assert eqIdent(forward[0], "forward")
  forward[^1].expectKind(nnkStmtList)
  # forward x:
  #   x.cv1.relu.mp1.flatten.classifier
  # Command
  #   Ident "forward"
  #   Ident "x"
  #   StmtList
  #     DotExpr
  #       DotExpr
  #         DotExpr
  #           DotExpr
  #             DotExpr
  #               Ident "x"
  #               Ident "cv1"
  #             Ident "relu"
  #           Ident "mp1"
  #         Ident "flatten"
  #       Ident "classifier"

  # 0. Prepare type information and the raw proc body
  let
    ModelType = newIdentNode(model_name)
    InOutType = nnkBracketExpr.newTree(
      newIdentNode("Variable"), self.subtype
    )
    body = forward[^1]

  # 1. Create the input variables with their type
  var inputVars = nnkIdentDefs.newTree()

  for varIndex in 1..forward.len-2:
    inputVars.add newIdentNode($forward[varIndex])

  inputVars.add InOutType
  inputVars.add newEmptyNode() # Default Value

  # 2. Add the shortut syntax templates
  var shortcutTemplates = newStmtList()

  for shortcut in self.forward_templates:
    shortcutTemplates.add shortcut

  # 3. Create the forward proc
  self.forward_proc = nnkProcDef.newTree(
    newIdentNode("forward"), newEmptyNode(), newEmptyNode(),
    nnkFormalParams.newTree(
      # Result type
      InOutType,
      # Model
      nnkIdentDefs.newTree(newIdentNode("self"), ModelType, newEmptyNode()),
      # Variables
      inputVars
    ),
    newEmptyNode(), newEmptyNode(),
    nnkStmtlist.newTree(
      # TODO asserts
      shortcutTemplates,
    )
  )

macro ctxSubtype(context: Context): untyped =
  ## Extract the subtype from a Context
  result = replaceSymsByIdents(context.getTypeInst[1])

macro network(ctx: Context, model_name: untyped, config: untyped): untyped =

  # 0. - Separate the configuration into layers and forward part
  #    - get the subtype of the model (Tensor[float32], CudaTensor[float64], ...)
  let sections = config.splitSections

  # 1. Initialize the VM to analyse the neural network Graph.
  #    - Get the input shapes
  #    - Get the layers
  let vm = new Neuromancer
  vm.context = ctx
  vm.subtype = getAST(ctxSubtype(ctx))
  vm.topoTable = initTable[NimNode, LayerTopology]()
  vm.topoTable.topoFromLayers(sections.layers)

  # 2. Generate the model fields, initialization and template synctactic sugar
  vm.genModelFieldInit()
  vm.genTemplateShortcuts()

  # 3. Generate the type section
  vm.genModelType($model_name)
  vm.genInitProc($model_name)
  vm.genForwardProc($model_name, sections.forward)

  # 4. Output the result: type + init proc + forward proc
  result = newStmtList()
  result.add vm.type_section
  result.add vm.init_proc
  result.add vm.forward_proc

var ctx: Context[Tensor[float32]]


network ctx, FooNet:
  layers:
    x:   Input([1, 28, 28])                            # Real shape [N, 1, 28, 28]
    cv1: Conv2D(x.out_shape, 20, 5, 5)                 # Output shape [N, 20, 24, 24] (kernel 5x5, padding 0, stride 1)
    mp1: MaxPool2D(cv1.out_shape, (2,2), (0,0), (2,2)) # Output shape [N, 20, 12, 12] (kernel 2X2, padding 0, stride 2)
    classifier:
      Linear(mp1.out_shape.flatten, 10)                # Output shape [N, 10]
  forward x:
    x.cv1.relu.mp1.flatten.classifier


let a = init(ctx, FooNet)


echo a