# batch_norm
out = g.op("BatchNormalization", input, weight, bias, running_mean, running_var,
            epsilon_f=eps,
            momentum_f=1 - momentum,
            outputs=1 if not training else 5)

# instance_norm
return g.op("InstanceNormalization", input, weight, bias, epsilon_f=eps)

BatchNormalization
{
    "Attr": [spatial, epsilon, momentum] # spatial这个属性一般不用 还有is_test这个属性在早期的ONNX版本也会出现
    "Input": [X, scale, B, mean, var]
    "Output": [Y, mean, var, save_mean, save_var]  #  后4个可选
}
InstanceNormalization
{
    "Attr": [epsilon, ]
    "Input": [input, scale, B]
    "Output": [output, ]
}

class InstanceNormalization(Base):

    @staticmethod
    def export():  # type: () -> None
        def _instancenorm_test_mode(x, s, bias, epsilon=1e-5):  # type: ignore
            dims_x = len(x.shape)
            axis = tuple(range(2, dims_x))
            mean = np.mean(x, axis=axis, keepdims=True)
            var = np.var(x, axis=axis, keepdims=True)
            dim_ones = (1,) * (dims_x - 2)
            s = s.reshape(-1, *dim_ones)
            bias = bias.reshape(-1, *dim_ones)
            return s * (x - mean) / np.sqrt(var + epsilon) + bias

sym, params = nnvm.frontend.from_ONNX(ONNX_model)
graph, lib, params = nnvm.compiler.build(sym, target, shape_dict, params=params, dtype={})

class BatchNorm(OnnxOpConverter):
    
    @classmethod
    def _impl_v1(cls, inputs, attr, params):
    # TODO(zhreshold): 'spatial' is not properly handled here.
        return AttrCvt(
            op_name='batch_norm',
            disables=['momentum'],
            ignores=['spatial', 'is_test', 'consumed_inputs'])(inputs, attr,
                                                                params)

class InstanceNorm(OnnxOpConverter):
    
    @classmethod
    def _impl_v1(cls, inputs, attr, params):
        # NOTE instance norm parameters is different from batch norm
        return AttrCvt(op_name='instance_norm')(inputs, attr, params)

@reg.register_compute("instance_norm")
def compute_instance_norm(attrs, inputs, _):
    eps = attrs.get_float("epsilon")
    return topi.nn.instance_norm_inference(inputs[0], inputs[1], inputs[2], eps, False)

@reg.register_schedule("instance_norm")
def schedule_instance_norm(attrs, outs, target):
    with tvm.target.create(target):
        return topi.generic.schedule_instance_norm(outs)

reg.register_pattern("instance_norm", OpPattern.OPAQUE)

def dense_default(data, weight, bias=None):
    assert len(data.shape) == 2 and len(weight.shape) == 2, \
        "only support 2-dim dense"
    if bias is not None:
        assert len(bias.shape) == 1
    batch, in_dim = data.shape
    out_dim, _ = weight.shape
    # 定义一个IterVar
    k = tvm.reduce_axis((0, in_dim), name='k')
    # 定义计算方式，tvm.compute的第一个参数为输出的形状，第二个参数为lambda表达式
    # lambda表达式中定义了输出中每个“坐标”处的值是如何计算的，其中有用到上面定义的IterVar
    matmul = tvm.compute((batch, out_dim), \
                         lambda i, j: tvm.sum(data[i, k] * weight[j, k], axis=k), \
                         tag='dense')
    if bias is not None:
        matmul = tvm.compute((batch, out_dim), \
                             lambda i, j: matmul[i, j] + bias[j], \
                             tag=tag.BROADCAST)
    return matmul

def instance_norm_inference(data, gamma, beta, eps, fix_gamma):
    assert len(data.shape) == 4, "only support 4-dim instance norm"
    batch, channel, height, width = data.shape
    rh = tvm.reduce_axis((0, height), name="rh")
    rw = tvm.reduce_axis((0, width), name="rw")
    # 当前batch, channel的均值
    s_mean = tvm.compute((batch, channel), \
            lambda b, c: tvm.sum(data[b, c, rh, rw] / (height * width), axis=[rh, rw]))
    # 再次使用相同的IterVar的时候需要重新定义一下，否则直接使用会报一个越界的错误
    rh = tvm.reduce_axis((0, height), name="rh")
    rw = tvm.reduce_axis((0, width), name="rw")
    s_mean_2 = tvm.compute((batch, channel), \
            lambda b, c: tvm.sum(tvm.intrin.power(data[b, c, rh, rw], 2) / (height * width), axis=[rh, rw]))
    # 当前batch, channel的方差
    s_var = tvm.compute((batch, channel), \
            lambda b, c: s_mean_2[b, c] - tvm.intrin.power(s_mean[b, c], 2))
    if fix_gamma:
        out = tvm.compute((batch, channel, height, width), \
            lambda b, c, h, w: (data[b, c, h, w] - s_mean[b, c]) / \
            tvm.intrin.sqrt(s_var[b, c] + eps) + beta[c])
    else:
        out = tvm.compute((batch, channel, height, width), \
            lambda b, c, h, w: (data[b, c, h, w] - s_mean[b, c]) / \
            tvm.intrin.sqrt(s_var[b, c] + eps) * gamma[c] + beta[c])
    return out

@generic.schedule_instance_norm.register(["cuda"])
def schedule_instance_norm(outs):
    # 从刚才的TOPI算法的最后返回tensor中，取出涉及到tvm.compute的几个tensor
    fin = outs[0].op.input_tensors
    mean = fin[1]
    var = fin[2]
    mean_2 = var.op.input_tensors[0]
    # instance_norm的计算过程中涉及到4个tvm.compute
    # 每个tvm.compute都需要Schedule
    ts = [mean, mean_2, var, outs[0]]

    s = tvm.create_schedule([x.op for x in ts])

    # 这里采用比较简单粗暴的方式，即每个tvm.compute都采用相同的Schedule方式
    # 事实上深入思考计算过程，还有进一步优化的可能，我也在做进一步的探索和尝试
    def _schedule(Out):  # the type of Out is tensor
        num_thread = 8
        block_x = tvm.thread_axis("blockIdx.x")
        block_y = tvm.thread_axis("blockIdx.y")
        thread_x = tvm.thread_axis((0, num_thread), "threadIdx.x")
        thread_y = tvm.thread_axis((0, num_thread), "threadIdx.y")

        by, ty = s[Out].split(s[Out].op.axis[0], factor=num_thread)
        bx, tx = s[Out].split(s[Out].op.axis[1], factor=num_thread)
        s[Out].reorder(by, bx, ty, tx)
        s[Out].bind(ty, thread_y)
        s[Out].bind(tx, thread_x)
        s[Out].bind(by, block_y)
        s[Out].bind(bx, block_x)

    for x in ts:
        _schedule(x)

    return s

def matmul(g, self, other):
    # return g.op("MatMul", self, other)
    # 上一行为原代码，下面为改动后的实现，其中C就是全0的bias
    ty = self.type().scalarType().lower()
    C = g.constant(0, [other.type().sizes()[1]], ty)
    # 这里稍微解释下GEMM的计算方式，假设输入有三个矩阵A,B,C，计算公式如下(这里没考虑有转置属性的情况)
    # alpha * AB + beta * C
    return g.op("Gemm", self, other, C, beta_f=0.0, alpha_f=1.0, broadcast_i=True)