# 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: E712, F401, F841

import json

import sys

import numpy as np

import pytest

import tvm_ffi

import tvm

import tvm.testing

from tvm import te

def test_const_saveload_json():

# save load json

x = tvm.tirx.const(1, "int32")

y = tvm.tirx.const(10, "int32")

z = x + y

z = z + z

json_str = tvm.ir.save_json(z)

zz = tvm.ir.load_json(json_str)

tvm.ir.assert_structural_equal(zz, z, map_free_vars=True)

def test_save_json_metadata_version():

obj = tvm.runtime.convert([1, 2])

json_str = tvm.ir.save_json(obj)

assert json.loads(json_str)["metadata"]["tvm_version"] == tvm.__version__

assert list(tvm.ir.load_json(json_str)) == [1, 2]

def _test_infinity_value(value, dtype):

x = tvm.tirx.const(value, dtype)

json_str = tvm.ir.save_json(x)

tvm.ir.assert_structural_equal(x, tvm.ir.load_json(json_str))

def test_infinity_value():

_test_infinity_value(float("inf"), "float64")

_test_infinity_value(float("-inf"), "float64")

_test_infinity_value(float("inf"), "float32")

_test_infinity_value(float("-inf"), "float32")

def _test_minmax_value(value):

json_str = tvm.ir.save_json(value)

tvm.ir.assert_structural_equal(value, tvm.ir.load_json(json_str))

def test_minmax_value():

_test_minmax_value(tvm.tirx.min_value("float32"))

_test_minmax_value(tvm.tirx.max_value("float32"))

def test_make_smap():

# save load json

x = tvm.tirx.const(1, "int32")

y = tvm.tirx.const(10, "int32")

z = tvm.tirx.Add(x, y)

smap = tvm.runtime.convert({"z": z, "x": x})

json_str = tvm.ir.save_json(tvm.runtime.convert([smap]))

arr = tvm.ir.load_json(json_str)

assert len(arr) == 1

assert arr[0]["z"].a == arr[0]["x"]

tvm.ir.assert_structural_equal(arr, [smap], map_free_vars=True)

def test_make_node():

x = tvm.ir.make_node("ir.IntImm", ty=tvm.ir.PrimType("int32"), value=10, span=None)

assert isinstance(x, tvm.tirx.IntImm)

assert x.value == 10

A = te.placeholder((10,), name="A")

AA = tvm.ir.make_node(

"te.Tensor", shape=A.shape, dtype=A.dtype, op=A.op, value_index=A.value_index

)

assert AA.op == A.op

assert AA.value_index == A.value_index

y = tvm.ir.make_node(

"ir.IntImm", ty=tvm.ir.PrimType(tvm_ffi.core.String("int32")), value=10, span=None

)

assert isinstance(y, tvm.tirx.IntImm)

assert y.value == 10

def test_make_sum():

A = te.placeholder((2, 10), name="A")

k = te.reduce_axis((0, 10), "k")

B = te.compute((2,), lambda i: te.sum(A[i, k], axis=k), name="B")

json_str = tvm.ir.save_json(B)

BB = tvm.ir.load_json(json_str)

assert B.op.body[0].combiner is not None

assert BB.op.body[0].combiner is not None

def test_string():

# non printable str, need to store by b64

s1 = tvm_ffi.core.String("xy\x01z")

s2 = tvm.ir.load_json(tvm.ir.save_json(s1))

tvm.ir.assert_structural_equal(s1, s2)

# printable str, need to store by repr_str

s1 = tvm_ffi.core.String("xyz")

s2 = tvm.ir.load_json(tvm.ir.save_json(s1))

tvm.ir.assert_structural_equal(s1, s2)

def test_pass_config():

cfg = tvm.transform.PassContext(

opt_level=1,

config={

"tirx.UnrollLoop": {

"auto_max_step": 10,

}

},

)

cfg.opt_level == 1

assert cfg.config["tirx.UnrollLoop"].auto_max_step == 10

# default option

assert cfg.config["tirx.UnrollLoop"].explicit_unroll == True

# schema checking for specific config key

with pytest.raises(TypeError):

cfg = tvm.transform.PassContext(config={"tirx.UnrollLoop": {"invalid": 1}})

# schema check for un-registered config

with pytest.raises(AttributeError):

cfg = tvm.transform.PassContext(config={"inavlid-opt": True})

# schema check for wrong type

with pytest.raises(AttributeError):

cfg = tvm.transform.PassContext(config={"tirx.UnrollLoop": 1})

def test_dict():

x = tvm.tirx.const(1) # a class that has Python-defined methods

# instances should see the full class dict

assert set(dir(x.__class__)) <= set(dir(x))

def test_tensor():

dev = tvm.cpu(0)

tvm_arr = tvm.runtime.tensor(np.random.rand(4), device=dev)

tvm_arr2 = tvm.ir.load_json(tvm.ir.save_json(tvm_arr))

tvm.ir.assert_structural_equal(tvm_arr, tvm_arr2)

np.testing.assert_array_equal(tvm_arr.numpy(), tvm_arr2.numpy())

def test_tensor_dict():

dev = tvm.cpu(0)

m1 = {

"key1": tvm.runtime.tensor(np.random.rand(4), device=dev),

"key2": tvm.runtime.tensor(np.random.rand(4), device=dev),

}

m2 = tvm.ir.load_json(tvm.ir.save_json(m1))

tvm.ir.assert_structural_equal(m1, m2)

def test_free_var_equal():

x = tvm.tirx.Var("x", dtype="int32")

y = tvm.tirx.Var("y", dtype="int32")

z = tvm.tirx.Var("z", dtype="int32")

v1 = x + y

v1 = y + z

tvm.ir.assert_structural_equal(x, z, map_free_vars=True)

if __name__ == "__main__":

tvm.testing.main()