# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: E402

"""
.. _relax-creation:

Relax Creation
==============
This tutorial demonstrates how to create Relax functions and programs.
We'll cover various ways to define Relax functions, including using TVMScript,
and relax NNModule API.
"""


######################################################################
# Create Relax programs using TVMScript
# -------------------------------------
# TVMScript is a domain-specific language for representing Apache TVM's
# intermediate representation (IR). It is a Python dialect that can be used
# to define an IRModule, which contains both TensorIR and Relax functions.
#
# In this section, we will show how to define a simple MLP model with only
# high-level Relax operators using TVMScript.

from tvm import relax, topi
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T


@I.ir_module
class RelaxModule:
    @R.function
    def forward(
        data: R.Tensor(("n", 784), dtype="float32"),
        w0: R.Tensor((128, 784), dtype="float32"),
        b0: R.Tensor((128,), dtype="float32"),
        w1: R.Tensor((10, 128), dtype="float32"),
        b1: R.Tensor((10,), dtype="float32"),
    ) -> R.Tensor(("n", 10), dtype="float32"):
        with R.dataflow():
            lv0 = R.matmul(data, R.permute_dims(w0)) + b0
            lv1 = R.nn.relu(lv0)
            lv2 = R.matmul(lv1, R.permute_dims(w1)) + b1
            R.output(lv2)
        return lv2


RelaxModule.show()

######################################################################
# Relax is not only a graph-level IR, but also supports cross-level
# representation and transformation. To be specific, we can directly call
# TensorIR functions in Relax function.


@I.ir_module
class RelaxModuleWithTIR:
    @T.prim_func(s_tir=True)
    def relu(x: T.handle, y: T.handle):
        n = T.int64()
        m = T.int64()
        X = T.match_buffer(x, (n, m), "float32")
        Y = T.match_buffer(y, (n, m), "float32")
        for i, j in T.grid(n, m):
            with T.sblock("relu"):
                vi, vj = T.axis.remap("SS", [i, j])
                Y[vi, vj] = T.max(X[vi, vj], T.float32(0))

    @R.function
    def forward(
        data: R.Tensor(("n", 784), dtype="float32"),
        w0: R.Tensor((128, 784), dtype="float32"),
        b0: R.Tensor((128,), dtype="float32"),
        w1: R.Tensor((10, 128), dtype="float32"),
        b1: R.Tensor((10,), dtype="float32"),
    ) -> R.Tensor(("n", 10), dtype="float32"):
        n = T.int64()
        cls = RelaxModuleWithTIR
        with R.dataflow():
            lv0 = R.matmul(data, R.permute_dims(w0)) + b0
            lv1 = R.call_tir(cls.relu, lv0, R.Tensor((n, 128), dtype="float32"))
            lv2 = R.matmul(lv1, R.permute_dims(w1)) + b1
            R.output(lv2)
        return lv2


RelaxModuleWithTIR.show()

######################################################################
# .. note::
#
#   You may notice that the printed output is different from the written
#   TVMScript code. This is because we print the IRModule in a standard
#   format, while we support syntax sugar for the input
#
#   For example, we can combine multiple operators into a single line, as
#
#   .. code-block:: python
#
#     lv0 = R.matmul(data, R.permute_dims(w0)) + b0
#
#   However, the normalized expression requires only one operation in one
#   binding. So the printed output is different from the written TVMScript code,
#   as
#
#   .. code-block:: python
#
#     lv: R.Tensor((784, 128), dtype="float32") = R.permute_dims(w0, axes=None)
#     lv1: R.Tensor((n, 128), dtype="float32") = R.matmul(data, lv)
#     lv0: R.Tensor((n, 128), dtype="float32") = R.add(lv1, b0)
#

######################################################################
# Create Relax programs using NNModule API
# ----------------------------------------
# Besides TVMScript, we also provide a PyTorch-like API for defining neural networks.
# It is designed to be more intuitive and easier to use than TVMScript.
#
# In this section, we will show how to define the same MLP model using
# Relax NNModule API.

from tvm.relax.frontend import nn


class NNModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu1(x)
        x = self.fc2(x)
        return x


######################################################################
# After we define the NNModule, we can export it to TVM IRModule via
# ``export_tvm``.

mod, params = NNModule().export_tvm({"forward": {"x": nn.spec.Tensor(("n", 784), "float32")}})
mod.show()

######################################################################
# We can also insert customized function calls into the NNModule, such as
# Tensor Expression(TE), TensorIR functions or other TVM packed functions.


@T.prim_func(s_tir=True)
def tir_linear(x: T.handle, w: T.handle, b: T.handle, z: T.handle):
    M = T.int64()
    N = T.int64()
    K = T.int64()
    X = T.match_buffer(x, (M, K), "float32")
    W = T.match_buffer(w, (N, K), "float32")
    B = T.match_buffer(b, (N,), "float32")
    Z = T.match_buffer(z, (M, N), "float32")
    for i, j, k in T.grid(M, N, K):
        with T.sblock("linear"):
            vi, vj, vk = T.axis.remap("SSR", [i, j, k])
            with T.init():
                Z[vi, vj] = 0
            Z[vi, vj] = Z[vi, vj] + X[vi, vk] * W[vj, vk]
    for i, j in T.grid(M, N):
        with T.sblock("add"):
            vi, vj = T.axis.remap("SS", [i, j])
            Z[vi, vj] = Z[vi, vj] + B[vj]


class NNModuleWithTIR(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        n = x.shape[0]
        # We can call external functions using nn.extern
        x = nn.extern(
            "env.linear",
            [x, self.fc1.weight, self.fc1.bias],
            out=nn.Tensor.placeholder((n, 128), "float32"),
        )
        # We can also call TensorIR via Tensor Expression API in TOPI
        x = nn.tensor_expr_op(topi.nn.relu, "relu", [x])
        # We can also call other TVM packed functions
        x = nn.tensor_ir_op(
            tir_linear,
            "tir_linear",
            [x, self.fc2.weight, self.fc2.bias],
            out=nn.Tensor.placeholder((n, 10), "float32"),
        )
        return x


mod, params = NNModuleWithTIR().export_tvm(
    {"forward": {"x": nn.spec.Tensor(("n", 784), "float32")}}
)
mod.show()


######################################################################
# Create Relax programs using Block Builder API
# ---------------------------------------------
# In addition to the above APIs, we also provide a Block Builder API for
# creating Relax programs. It is a IR builder API, which is more
# low-level and widely used in TVM's internal logic, e.g writing a
# customized pass.

bb = relax.BlockBuilder()
n = T.int64()
x = relax.Var("x", R.Tensor((n, 784), "float32"))
fc1_weight = relax.Var("fc1_weight", R.Tensor((128, 784), "float32"))
fc1_bias = relax.Var("fc1_bias", R.Tensor((128,), "float32"))
fc2_weight = relax.Var("fc2_weight", R.Tensor((10, 128), "float32"))
fc2_bias = relax.Var("fc2_bias", R.Tensor((10,), "float32"))
with bb.function("forward", [x, fc1_weight, fc1_bias, fc2_weight, fc2_bias]):
    with bb.dataflow():
        lv0 = bb.emit(relax.op.matmul(x, relax.op.permute_dims(fc1_weight)) + fc1_bias)
        lv1 = bb.emit(relax.op.nn.relu(lv0))
        gv = bb.emit(relax.op.matmul(lv1, relax.op.permute_dims(fc2_weight)) + fc2_bias)
        bb.emit_output(gv)
    bb.emit_func_output(gv)

mod = bb.get()
mod.show()

######################################################################
# Also, Block Builder API supports building cross-level IRModule with both
# Relax functions, TensorIR functions and other TVM packed functions.

bb = relax.BlockBuilder()
with bb.function("forward", [x, fc1_weight, fc1_bias, fc2_weight, fc2_bias]):
    with bb.dataflow():
        lv0 = bb.emit(
            relax.call_dps_packed(
                "env.linear",
                [x, fc1_weight, fc1_bias],
                out_ty=relax.TensorType((n, 128), "float32"),
            )
        )
        lv1 = bb.emit_te(topi.nn.relu, lv0)
        tir_gv = bb.add_func(tir_linear, "tir_linear")
        gv = bb.emit(
            relax.call_tir(
                tir_gv,
                [lv1, fc2_weight, fc2_bias],
                out_ty=relax.TensorType((n, 10), "float32"),
            )
        )
        bb.emit_output(gv)
    bb.emit_func_output(gv)
mod = bb.get()
mod.show()

######################################################################
# Note that the Block Builder API is not as user-friendly as the above APIs,
# but it is lowest-level API and works closely with the IR definition. We
# recommend using the above APIs for users who only want to define and
# transform a ML model. But for those who want to build more complex
# transformations, the Block Builder API is a more flexible choice.

######################################################################
# Summary
# -------
# This tutorial demonstrates how to create Relax programs using TVMScript,
# NNModule API, Block Builder API and PackedFunc API for different use cases.
