"""x86 nn operators"""
from __future__ import absolute_import as _abs
import tvm
from .. import generic

def _default_schedule(outs, auto_inline):
    """Default schedule for x86."""
    x = outs[0]
    s = tvm.create_schedule([x.op for x in outs])
    if auto_inline:
        tvm.schedule.AutoInlineInjective(s)
        s[x].fuse(s[x].op.axis)
        return s
    if len(s[x].op.axis) == 4:
        n, c, _, _ = s[x].op.axis
        fused = s[x].fuse(n, c) # for nhwc layout, fuse n and h
        s[x].parallel(fused)
    else:
        s[x].parallel(s[x].op.axis[0])
    return s


@generic.schedule_softmax.register(["cpu"])
def schedule_softmax(outs):
    """Schedule for softmax

    Parameters
    ----------
    outs: Array of Tensor
          The computation graph description of softmax
          in the format of an array of tensors.

    Returns
    -------
    sch: Schedule
        The computation schedule for the op.
    """
    return _default_schedule(outs, False)


@generic.schedule_pool.register(["cpu"])
def schedule_pool(outs):
    """Schedule for pool

    Parameters
    ----------
    outs: Array of Tensor
          The computation graph description of pool
          in the format of an array of tensors.

    Returns
    -------
    sch: Schedule
        The computation schedule for the op.
    """
    return _default_schedule(outs, False)

@@ -1,8 +1,55 @@
+# pylint: disable=invalid-name,too-many-locals%0A
 %22%22%22x86 n
@@ -143,16 +143,35 @@
  generic
+%0Afrom .. import tag
 %0A%0Adef _d
@@ -1379,32 +1379,32 @@
 the op.%0A    %22%22%22%0A
-
     return _defa
@@ -1409,28 +1409,1600 @@
 fault_schedule(outs, False)%0A
+%0A%0A@generic.schedule_dense.register(%5B%22cpu%22%5D)%0Adef schedule_dense(outs):%0A    %22%22%22Schedule for dense%0A%0A    Parameters%0A    ----------%0A    outs: Array of Tensor%0A          The computation graph description of pool%0A          in the format of an array of tensors.%0A%0A    Returns%0A    -------%0A    sch: Schedule%0A        The computation schedule for the op.%0A    %22%22%22%0A%0A    outs = %5Bouts%5D if isinstance(outs, tvm.tensor.Tensor) else outs%0A    s = tvm.create_schedule(%5Bx.op for x in outs%5D)%0A%0A    def traverse(op):%0A        %22%22%22Traverse operators from computation graph%22%22%22%0A        # inline all one-to-one-mapping operators except the last stage (output)%0A        if tag.is_broadcast(op.tag):%0A            if op not in s.outputs:%0A                s%5Bop%5D.compute_inline()%0A            for tensor in op.input_tensors:%0A                if tensor.op.input_tensors:%0A                    traverse(tensor.op)%0A%0A        if 'dense' in op.tag:%0A            C = op.output(0)%0A            x, y = C.op.axis%0A%0A            # Write cache for blocks%0A            CC = s.cache_write(C, 'global')%0A%0A            # Tile%0A            bnx = 1%0A            bny = 4%0A            _, yo, _, yi = s%5BC%5D.tile(x, y, bnx, bny)%0A            s%5BCC%5D.compute_at(s%5BC%5D, yo)%0A            xc, yc = s%5BCC%5D.op.axis%0A            k, = s%5BCC%5D.op.reduce_axis%0A            ko, ki = s%5BCC%5D.split(k, factor=4)%0A            s%5BCC%5D.reorder(ko, xc, ki, yc)%0A            s%5BCC%5D.unroll(ki)%0A            s%5BCC%5D.vectorize(yc)%0A%0A            # Vectorization%0A            s%5BC%5D.vectorize(yi)%0A%0A            # Parallelization%0A            s%5BC%5D.parallel(yo)%0A%0A    traverse(outs%5B0%5D.op)%0A    return s%0A

Add dense base scheduler (#887)