text: # 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.
  import pytest
  import tvm
  from tvm import te
  import numpy as np
  
  
  def test_const():
      x = tvm.tir.const(1, "int32")
      assert x.dtype == "int32"
      assert isinstance(x, tvm.tir.IntImm)
  
  
  def test_te_const():
      x = tvm.te.const(1, "int32")
      assert x.dtype == "int32"
      assert isinstance(x, tvm.tir.IntImm)
  
  
  def test_scalar_dtype_inference():
      for data in [
          True,
          bool(1),
          np.uint8(1),
          np.uint16(1),
          np.uint32(1),
          np.uint64(1),
          np.int8(1),
          np.int16(1),
          np.int32(1),
          np.int64(1),
          np.float16(1),
          np.float32(1),
          np.float64(1),
      ]:
          assert tvm.tir.const(data).dtype == str(np.array(data).dtype)
      assert tvm.tir.const(1).dtype == "int32"
      assert tvm.tir.const(1.0).dtype == "float32"
  
      for data in [
          True,
          bool(1),
          np.uint8(1),
          np.uint16(1),
          np.uint32(1),
          np.uint64(1),
          np.int8(1),
          np.int16(1),
          np.int32(1),
          np.int64(1),
          np.float16(1),
          np.float32(1),
          np.float64(1),
      ]:
          assert tvm.runtime.convert(data).dtype == str(np.array(data).dtype)
      assert tvm.runtime.convert(1).dtype == "int32"
      assert tvm.runtime.convert(1.0).dtype == "float32"
  
  
  def test_make():
      x = tvm.tir.const(1, "int32")
      y = te.var("x")
      z = x + y
      assert isinstance(tvm.tir.max(x, y), tvm.tir.Max)
      assert isinstance(tvm.tir.min(x, y), tvm.tir.Min)
  
  
  def test_ir():
      x = tvm.tir.const(1, "int32")
      y = tvm.tir.IntImm("int32", 1)
      z = x + y
      stmt = tvm.tir.Evaluate(z)
      assert isinstance(stmt, tvm.tir.Evaluate)
  
  
  def test_ir2():
      x = te.var("n")
      a = te.var("array", "handle")
      st = tvm.tir.Store(a, x + 1, 1)
      assert isinstance(st, tvm.tir.Store)
      assert st.buffer_var == a
  
  
  def test_let():
      x = te.var("x")
      y = te.var("y")
      stmt = tvm.tir.LetStmt(x, 10, tvm.tir.Evaluate(x + 1))
  
  
  def test_cast():
      x = te.var("x", dtype="float32")
      y = x.astype("int32")
      z = x.astype("float32x4")
      assert isinstance(y, tvm.tir.Cast)
      assert isinstance(z, tvm.tir.Broadcast)
      assert z.lanes == 4
  
      s = tvm.tir.StringImm("s")
      with pytest.raises(tvm.error.TVMError) as cm:
          s.astype("int")
          assert "Can't cast a handle to other types" in str(cm.execption)
  
  
  def test_attr():
      x = te.var("x")
      y = te.var("y")
      stmt = tvm.tir.AttrStmt(y, "stride", 10, tvm.tir.Evaluate(x + 1))
      assert stmt.node == y
  
      a = tvm.runtime.convert(1)
      assert a.value == 1
      try:
          a.no_field
          assert False
      except AttributeError:
          pass
  
  
  def test_basic():
      a = te.var("a")
      b = te.var("b")
      c = a + b
      assert str(c) == "(%s: int32 + %s: int32)" % (a.name, b.name)
  
  
  def test_stmt():
      x = tvm.tir.Evaluate(0)
      tvm.tir.For(te.var("i"), 0, 1, tvm.tir.ForKind.SERIAL, x)
  
  
  def test_dir():
      x = te.var("x")
      dir(x)
  
  
  def test_dtype():
      x = te.var("x")
      assert x.dtype == "int32"
      y = te.var("y")
      assert (x > y).dtype == "bool"
  
  
  def test_any():
      x = te.var("x")
      y = te.var("y")
      z = te.var("z")
      try:
          t = x or x
          assert False
      except ValueError:
          pass
      try:
          tvm.tir.any()
          assert False
      except ValueError:
          pass
      assert str(tvm.tir.any(x < y)) == "(%s: int32 < %s: int32)" % (x.name, y.name)
      assert str(tvm.tir.any(x < y, x > z)) == "((%s: int32 < %s: int32) || (%s > %s: int32))" % (
          x.name,
          y.name,
          x.name,
          z.name,
      )
      assert str(
          tvm.tir.any(x < y, y > z + 1, x < z * 2)
      ) == "(((%s: int32 < %s: int32) || (%s > (%s: int32 + 1))) || (%s < (%s*2)))" % (
          x.name,
          y.name,
          y.name,
          z.name,
          x.name,
          z.name,
      )
  
  
  def test_all():
      x = te.var("x")
      y = te.var("y")
      z = te.var("z")
      try:
          t = x and x
          assert False
      except ValueError:
          pass
      try:
          tvm.tir.all()
          assert False
      except ValueError:
          pass
      assert str(tvm.tir.all(x < y)) == "(%s: int32 < %s: int32)" % (x.name, y.name)
      assert str(tvm.tir.all(x < y, x > z)) == "((%s: int32 < %s: int32) && (%s > %s: int32))" % (
          x.name,
          y.name,
          x.name,
          z.name,
      )
      assert str(
          tvm.tir.all(x < y, y > z + 1, x < z * 2)
      ) == "(((%s: int32 < %s: int32) && (%s > (%s: int32 + 1))) && (%s < (%s*2)))" % (
          x.name,
          y.name,
          y.name,
          z.name,
          x.name,
          z.name,
      )
  
  
  def test_bitwise():
      x = te.var("x")
      y = te.var("y")
      assert str(x << y) == "@tir.shift_left(x: int32, y: int32, dtype=int32)"
      assert str(x >> y) == "@tir.shift_right(x: int32, y: int32, dtype=int32)"
      assert str(x & y) == "@tir.bitwise_and(x: int32, y: int32, dtype=int32)"
      assert str(x | y) == "@tir.bitwise_or(x: int32, y: int32, dtype=int32)"
      assert str(x ^ y) == "@tir.bitwise_xor(x: int32, y: int32, dtype=int32)"
      assert str(10 & x) == "@tir.bitwise_and(10, x: int32, dtype=int32)"
      assert str(10 | x) == "@tir.bitwise_or(10, x: int32, dtype=int32)"
      assert str(10 ^ x) == "@tir.bitwise_xor(10, x: int32, dtype=int32)"
      assert str(10 >> x) == "@tir.shift_right(10, x: int32, dtype=int32)"
      assert str(10 << x) == "@tir.shift_left(10, x: int32, dtype=int32)"
      assert str(10 % x) == "floormod(10, x: int32)"
  
      assert str(~x) == "@tir.bitwise_not(x: int32, dtype=int32)"
      assert (tvm.tir.const(1, "int8x2") >> 1).dtype == "int8x2"
      assert (x >> tvm.tir.const(1, "int32x2")).dtype == "int32x2"
      assert (te.var("z", "int8x2") << tvm.tir.const(1, "int8x2")).dtype == "int8x2"
  
  
  def test_float_bitwise():
      t = tvm.tir.const(1.5, dtype="float32")
      for test in [
          lambda lhs, rhs: lhs << rhs,
          lambda lhs, rhs: lhs >> rhs,
          lambda lhs, rhs: lhs | rhs,
          lambda lhs, rhs: lhs ^ rhs,
          lambda lhs, rhs: lhs & rhs,
      ]:
          try:
              test(t, 10.0)
              assert False
          except tvm.TVMError:
              pass
      try:
          ~t
          assert False
      except RuntimeError:
          pass
  
  
  def test_shift_bounds():
      x = te.var("x")
      for test in [lambda lhs, rhs: lhs << rhs, lambda lhs, rhs: lhs >> rhs]:
          # negative case
          for testcase in [(x, -1), (x, 32)]:
              try:
                  test(*testcase)
                  assert False
              except tvm.TVMError:
                  pass
  
          # positive case
          for testcase in [(x, 0), (x, 16), (x, 31)]:
              test(*testcase)
  
  
  def test_divide_by_zero():
      for test in [
          lambda lhs, rhs: tvm.tir.floormod(lhs, rhs),
          lambda lhs, rhs: tvm.tir.floordiv(lhs, rhs),
          lambda lhs, rhs: tvm.tir.truncmod(lhs, rhs),
          lambda lhs, rhs: tvm.tir.truncdiv(lhs, rhs),
          lambda lhs, rhs: tvm.tir.div(lhs, rhs),
      ]:
          try:
              test(tvm.tir.const(5, "int32"), tvm.tir.const(0, "int32"))
              assert False
          except tvm.TVMError:
              pass
  
  
  def test_isnan():
      x = te.var("x", "float32")
      assert str(tvm.tir.isnan(x)) == "@tir.isnan(x: float32, dtype=bool)"
      assert str(tvm.tir.isnan(x).dtype) == "bool"
      y = te.var("y", "float16")
      assert str(tvm.tir.isnan(y)) == "@tir.isnan(cast(float32, y: float16), dtype=bool)"
      z = te.var("z", "int32")
      assert str(tvm.tir.isnan(z)) == "False"
      k = te.var("k", "int8x2")
      assert str(tvm.tir.isnan(k).dtype) == "uint1x2"
  
  
  def test_equality():
      a = te.var("a")
      b = te.var("b")
      c = a == b
      assert not c
      d = c != c
      assert not d
  
  
  def test_equality_string_imm():
      x = "a"
      y = tvm.tir.StringImm(x)
      x == y.value
      x == y
  
  
  def test_prim_func():
      x = te.var("x")
      y = te.var("y")
      b = tvm.tir.decl_buffer((x,), "float32")
      stmt = tvm.tir.LetStmt(x, 10, tvm.tir.Evaluate(x + 1))
  
      func = tvm.tir.PrimFunc([x, y, b], stmt)
      # make sure we can print
      func.astext()
      assert func.buffer_map[func.params[2]].same_as(b)
  
      assert len(func.buffer_map) == 1
      f2 = func.with_attr({"calling_conv": 1, "tir.noalias": True})
      assert f2.attrs["calling_conv"].value == 1
      assert func.attrs is None
  
  
  def test_vars():
      x = tvm.tir.Var("xyz", "int8")
      assert x.dtype == "int8"
      ptype = tvm.ir.PointerType(tvm.ir.PrimType("float"))
      x = tvm.tir.Var("xyz", ptype)
      assert x.dtype == "handle"
      assert x.type_annotation == ptype
      assert isinstance(ptype.element_type, tvm.ir.PrimType)
  
  
  def test_scoped_storage_vars():
      dtype = "float"
      storage_scope = "global.texture"
      ptype = tvm.ir.PointerType(tvm.ir.PrimType(dtype), storage_scope)
      x = tvm.tir.Var("xyz", ptype)
      assert x.dtype == "handle"
      assert x.type_annotation == ptype
      assert x.type_annotation.storage_scope == storage_scope
      assert isinstance(ptype.element_type, tvm.ir.PrimType)
  
  
  def test_buffer_load_store():
      b = tvm.tir.decl_buffer((10,), "float32")
      x = tvm.tir.BufferLoad(b, [0])
      assert isinstance(x, tvm.tir.BufferLoad)
      assert x.dtype == "float32"
      assert x.buffer == b
      s = tvm.tir.BufferStore(b, 0.1, [0])
      assert isinstance(s, tvm.tir.BufferStore)
  
      s = tvm.tir.BufferRealize(b, [tvm.ir.Range(0, 1)], True, tvm.tir.Evaluate(0))
      assert isinstance(s, tvm.tir.BufferRealize)
  
  
  def test_intimm_cond():
      x = tvm.runtime.convert(1)
      y = tvm.runtime.convert(1)
      s = {x}
      assert y in s
      assert x == y
      assert x < 20
      assert not (x >= 20)
      assert x < 10 and y < 10
      assert not tvm.runtime.convert(x != 1)
      assert x == 1
  
  
  def test_block_blockrealize():
      x = tvm.tir.Var("x", "int32")
      y = tvm.tir.Var("y", "int32")
      vx = tvm.tir.IterVar((16, 16), "vx", 0)
      vx_var = vx.var
      vy = tvm.tir.IterVar((16, 16), "vy", 2)
      vy_var = vy.var
      A = tvm.tir.decl_buffer((16), "float32")
      B = tvm.tir.decl_buffer((16, 16), "float32")
      alloc_buffer = tvm.tir.decl_buffer((16, 16), "float32")
      match_buffer = tvm.tir.decl_buffer((16, 16), "float32")
      init_body = tvm.tir.BufferStore(A, 0.0, [vx_var])
      body = tvm.tir.BufferStore(
          A,
          tvm.tir.BufferLoad(A, [vx_var]) + tvm.tir.BufferLoad(B, [vx_var, vy_var]),
          [vx_var],
      )
      reads = [
          tvm.tir.BufferRegion(
              B, [tvm.ir.Range.from_min_extent(vx_var, 1), tvm.ir.Range.from_min_extent(vy_var, 1)]
          )
      ]
      writes = [tvm.tir.BufferRegion(A, [tvm.ir.Range.from_min_extent(vx_var, 1)])]
      block_match_buffer = tvm.tir.MatchBufferRegion(
          match_buffer, tvm.tir.BufferRegion(B, [tvm.ir.Range(0, 16), tvm.ir.Range(0, 16)])
      )
  
      block = tvm.tir.Block(
          [vx, vy],
          reads,
          writes,
          "block",
          body,
          init=init_body,
          alloc_buffers=[alloc_buffer],
          match_buffers=[block_match_buffer],
          annotations={"attr_key": "attr_value"},
      )
  
      # Checking Block
      assert isinstance(block, tvm.tir.Block)
      # Checking iter_vars
      assert block.iter_vars[0] == vx
      assert block.iter_vars[1] == vy
      # Checking reads/writes region
      assert isinstance(block.reads[0], tvm.tir.BufferRegion)
      assert block.reads[0].buffer == B
      assert block.reads[0].region[0].min == vx_var
      assert block.reads[0].region[1].min == vy_var
      assert isinstance(block.writes[0], tvm.tir.BufferRegion)
      assert block.writes[0].buffer == A
      assert block.writes[0].region[0].min == vx_var
      assert block.writes[0].region[0].extent == 1
      # Checking name_hint
      assert block.name_hint == "block"
      # Checking body
      assert block.body == body
      # Checking init
      assert block.init == init_body
      # Checking alloc_buffers
      assert block.alloc_buffers[0] == alloc_buffer
      # Checking match_buffers
      assert block.match_buffers[0].buffer == match_buffer
      assert isinstance(block.match_buffers[0].source, tvm.tir.BufferRegion)
      assert block.match_buffers[0].source.buffer == B
      assert block.match_buffers[0].source.region[0].min == 0
      assert block.match_buffers[0].source.region[0].extent == 16
  
      # Checking BlockRealize
      block_realize = tvm.tir.BlockRealize([x, y], tvm.tir.const(True, "bool"), block)
      assert isinstance(block_realize, tvm.tir.BlockRealize)
      assert block_realize.iter_values[0] == x
      assert block_realize.iter_values[1] == y
      assert block_realize.predicate == tvm.tir.const(True, "bool")
      assert block_realize.block == block
  
      # make sure we can print using ReprPrinter
      str(block)
      str(block_realize)
      # make sure we can print using TIRTextPrinter
      func = tvm.tir.PrimFunc([], block_realize)
      output = func.astext()
      assert output.find("meta[tir.BlockRealise]") == -1
      assert output.find("bind") != -1
      assert output.find("reads") != -1
      assert output.find("writes") != -1
      assert output.find("alloc_buffer") != -1
      assert output.find("match_buffer") != -1
      assert output.find("attr") != -1
      assert output.find("with init()") != -1
  
  
  def test_tir_allocate():
      dtype = "int8"
      storage_scope = "global"
      ptype = tvm.ir.PointerType(tvm.ir.PrimType(dtype), storage_scope)
      a = te.var("buffer", ptype)
      allocate = tvm.tir.Allocate(
          buffer_var=a,
          dtype=dtype,
          extents=[2, 2],
          condition=tvm.get_global_func("tir.const_true")(dtype, None),
          body=tvm.tir.Evaluate(2 + 1),
          annotations={
              "attr1": "foo",
              "attr2": "bar",
          },
      )
      assert allocate.buffer_var == a
      assert allocate.dtype == "int8"
      assert list(allocate.extents) == [2, 2]
      assert allocate.annotations["attr1"] == "foo"
      assert allocate.annotations["attr2"] == "bar"
  
      # make sure we can print using TIRTextPrinter
      func = tvm.tir.PrimFunc([], allocate)
      output = func.astext()
      assert (
          output.find(
              'allocate(buffer: Pointer(global int8), int8, [2, 2]), storage_scope = global, annotations = {"attr2": "bar", "attr1": "foo"})'
          )
          != -1
      )
  
  
  if __name__ == "__main__":
      pytest.main([__file__])
  