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 tvm_ext
  import tvm
  import tvm._ffi.registry
  import tvm.testing
  from tvm import te
  import numpy as np
  
  
  def test_bind_add():
      def add(a, b):
          return a + b
  
      f = tvm_ext.bind_add(add, 1)
      assert f(2) == 3
  
  
  def test_ext_dev():
      n = 10
      A = te.placeholder((n,), name="A")
      B = te.compute((n,), lambda *i: A(*i) + 1.0, name="B")
      s = te.create_schedule(B.op)
  
      def check_llvm():
          if not tvm.testing.device_enabled("llvm"):
              return
          f = tvm.build(s, [A, B], "ext_dev", "llvm")
          dev = tvm.ext_dev(0)
          # launch the kernel.
          a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)
          b = tvm.nd.array(np.zeros(n, dtype=B.dtype), dev)
          f(a, b)
          tvm.testing.assert_allclose(b.numpy(), a.numpy() + 1)
  
      check_llvm()
  
  
  def test_sym_add():
      a = te.var("a")
      b = te.var("b")
      c = tvm_ext.sym_add(a, b)
      assert c.a == a and c.b == b
  
  
  def test_ext_vec():
      ivec = tvm_ext.ivec_create(1, 2, 3)
      assert isinstance(ivec, tvm_ext.IntVec)
      assert ivec[0] == 1
      assert ivec[1] == 2
  
      def ivec_cb(v2):
          assert isinstance(v2, tvm_ext.IntVec)
          assert v2[2] == 3
  
      tvm.runtime.convert(ivec_cb)(ivec)
  
  
  def test_extract_ext():
      fdict = tvm._ffi.registry.extract_ext_funcs(tvm_ext._LIB.TVMExtDeclare)
      assert fdict["mul"](3, 4) == 12
  
  
  def test_extern_call():
      n = 10
      A = te.placeholder((n,), name="A")
      B = te.compute(
          (n,), lambda *i: tvm.tir.call_extern("float32", "TVMTestAddOne", A(*i)), name="B"
      )
      s = te.create_schedule(B.op)
  
      def check_llvm():
          if not tvm.testing.device_enabled("llvm"):
              return
          f = tvm.build(s, [A, B], "llvm")
          dev = tvm.cpu(0)
          # launch the kernel.
          a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)
          b = tvm.nd.array(np.zeros(n, dtype=B.dtype), dev)
          f(a, b)
          tvm.testing.assert_allclose(b.numpy(), a.numpy() + 1)
  
      check_llvm()
  
  
  def test_nd_subclass():
      a = tvm_ext.NDSubClass.create(additional_info=3)
      b = tvm_ext.NDSubClass.create(additional_info=5)
      assert isinstance(a, tvm_ext.NDSubClass)
      c = a + b
      d = a + a
      e = b + b
      assert a.additional_info == 3
      assert b.additional_info == 5
      assert c.additional_info == 8
      assert d.additional_info == 6
      assert e.additional_info == 10
  
  
  if __name__ == "__main__":
      test_nd_subclass()
      test_extern_call()
      test_ext_dev()
      test_ext_vec()
      test_bind_add()
      test_sym_add()
      test_extract_ext()
  