code: def test_in_bounds_const_loop_partition_ir():
      def check_attr_stmt(x):
          if (
              isinstance(x, tvm.tir.AttrStmt)
              and x.attr_key == "buffer_bound"
              and tvm.ir.structural_equal(x.value.args, [n])
          ):
              return True
          return False
  
      def check_branch_stmt(x):
          if isinstance(x, tvm.tir.IfThenElse):
              return True
          return False
  
      def assert_bound_instrumentation(stmt, f, nums):
          count = 0
          for i in collect_visit(stmt, f):
              if i is True:
                  count = count + 1
          assert count == nums
  
      def collect_branch_stmt(x):
          if isinstance(x, tvm.tir.IfThenElse):
              branch_collector.append(x)
  
      n = 21
      A = te.placeholder((n,), name="A")
      B = te.placeholder((n,), name="B")
  
      T = te.compute((n,), lambda i: A[i] + B[i])
      s = te.create_schedule(T.op)
      xo, xi = s[T].split(T.op.axis[0], factor=4)
  
      with tvm.transform.PassContext(
          config={
              "tir.instrument_bound_checkers": True,
              "tir.LoopPartition": {"partition_const_loop": True},
          }
      ):
          mod = tvm.driver.lower(s, [A, B, T], name="main")
  
      stmt = mod["main"].body
      # after instrumentation
      assert_bound_instrumentation(stmt, check_attr_stmt, 2 * 3)
      assert_bound_instrumentation(stmt, check_branch_stmt, 2)
  
      branch_collector = list()
      collect_visit(stmt, collect_branch_stmt)
      assert len(branch_collector) == 2
  
imports: import pytest
  import tvm
  import tvm.testing
  from tvm import te
  import numpy
label_desc: Use image node_id 8 for calling a global function with example usage: test_in_bounds_const_loop_partition_ir() and returns: False, False, True, True