code: def test_unroll_loop():
      ib = tvm.tir.ir_builder.create()
      dtype = "int64"
      n = te.size_var("n")
      Ab = tvm.tir.decl_buffer((n,), dtype)
      Aptr = ib.buffer_ptr(Ab)
      # for i in 0 to n-1:
      with ib.for_range(n, n + 2, name="i") as i:
          with ib.for_range(0, 8, name="i", kind="unroll") as j:
              Aptr[j + 1] = Aptr[i] + 1
  
      stmt = ib.get()
      mod = tvm.IRModule.from_expr(tvm.tir.PrimFunc([Ab], stmt))
  
      assert isinstance(stmt, tvm.tir.For)
  
      with tvm.transform.PassContext(
          config={"tir.UnrollLoop": {"auto_max_step": 16}}
      ):
          ret = tvm.tir.transform.UnrollLoop()(mod)["main"].body
          assert not isinstance(ret, tvm.tir.For)
  
      with tvm.transform.PassContext(
          config={"tir.UnrollLoop": {"auto_max_step": 15}}
      ):
          ret = tvm.tir.transform.UnrollLoop()(mod)["main"].body
          assert isinstance(ret, tvm.tir.For)
  
      with tvm.transform.PassContext(
          config={
              "tir.UnrollLoop": {
                  "auto_max_step": 16,
                  "explicit_unroll": False,
              }
          }
      ):
          ret = tvm.tir.transform.UnrollLoop()(mod)["main"].body
          assert isinstance(ret, tvm.tir.For)
          assert ret.kind == tvm.tir.ForKind.UNROLLED
  
      ib = tvm.tir.ir_builder.create()
      ib.scope_attr(
          tvm.tir.const(0, "int32"), "pragma_auto_unroll_max_step", 16
      )
      ib.emit(stmt)
      wrapped = ib.get()
      wrapped = tvm.tir.SeqStmt([wrapped, stmt])
      assert isinstance(ret, tvm.tir.For)
      mod = tvm.IRModule.from_expr(tvm.tir.PrimFunc([Ab], wrapped))
  
      with tvm.transform.PassContext(
          config={
              "tir.UnrollLoop": {
                  "auto_max_depth": 8,
                  "explicit_unroll": False,
              }
          }
      ):
          ret = tvm.tir.transform.UnrollLoop()(mod)["main"].body
          assert isinstance(ret[0], tvm.tir.For)
          assert ret[0].kind == tvm.tir.ForKind.UNROLLED
          assert isinstance(ret[1], tvm.tir.For)
          assert ret[1].kind != tvm.tir.ForKind.UNROLLED
  
file_path: tvm/tests/python/tir-transform/test_tir_transform_unroll_loop.py
imports: import tvm
  from tvm import te
  from tvm.script import tir
  import os
label_desc: Use image node_id 1 for calling a global function with example usage: test_unroll_loop() without return types