code: def FoldUopLoop():
      """Detect and fold uop loop.
  
      VTA support uop programming model
      that recognizes loop structure.
      This pass detect the loop structure
      and extract that into uop loop AST.
  
      Returns
      -------
      fpass : tvm.transform.Pass
          The pass
      """
  
      def _fold_outermost_loop(body):
          stmt = body
          if not isinstance(stmt, tvm.tir.For):
              return None, body, None
  
          loop_var = stmt.loop_var
          gemm_offsets = [None, None, None]
          fail = [False]
          builtin_uop_push = tvm.ir.Op.get("tir.vta.uop_push")
  
          def _post_order(op):
              assert isinstance(op, tvm.tir.Call)
              base_args = 2
              if op.op.same_as(builtin_uop_push):
                  args = []
                  args += op.args[:base_args]
                  for i in range(3):
                      m = tvm.arith.detect_linear_equation(
                          op.args[i + base_args], [loop_var]
                      )
                      if not m:
                          fail[0] = True
                          return op
                      if gemm_offsets[i] is not None:
                          if not tvm.ir.structural_equal(
                              m[0], gemm_offsets[i]
                          ):
                              fail[0] = True
                              return op
                          args.append(m[1])
                      else:
                          gemm_offsets[i] = m[0]
                          args.append(m[1])
                  args += op.args[base_args + 3 :]
                  return tvm.tir.call_intrin(
                      "int32", builtin_uop_push, *args
                  )
              if op.op.name not in (
                  "tir.vta.command_handle",
                  "tir.tvm_thread_context",
              ):
                  raise RuntimeError("unexpected op %s" % op)
              return op
  
          ret = tvm.tir.stmt_functor.ir_transform(
              stmt.body, None, _post_order, ["tir.Call"]
          )
  
          if not fail[0] and all(x is not None for x in gemm_offsets):
  
              def _visit(op):
                  if op.same_as(loop_var):
                      fail[0] = True
  
              tvm.tir.stmt_functor.post_order_visit(ret, _visit)
              if not fail[0]:
                  begin = tvm.tir.call_extern(
                      "int32",
                      "VTAUopLoopBegin",
                      stmt.extent,
                      *gemm_offsets,
                  )
                  end = tvm.tir.call_extern("int32", "VTAUopLoopEnd")
                  return [begin, ret, end]
          raise ValueError("Failed to fold the GEMM instructions..")
  
      def _do_fold(stmt):
          env = get_env()
          if (
              stmt.attr_key == "coproc_uop_scope"
              and isinstance(stmt.value, tvm.tir.StringImm)
              and stmt.value.value == env.dev.vta_push_uop.value
          ):
              body = stmt.body
              begins = []
              ends = []
              try:
                  begin, body, end = _fold_outermost_loop(body)
                  if begin is not None:
                      begins.append(begin)
                  if end is not None:
                      ends.append(end)
                  begin, body, end = _fold_outermost_loop(body)
                  if begin is not None:
                      begins.append(begin)
                  if end is not None:
                      ends.append(end)
              except ValueError:
                  pass
              if body == stmt.body:
                  return stmt
              ends = list(reversed(ends))
              body = tvm.tir.stmt_seq(*(begins + [body] + ends))
              return tvm.tir.AttrStmt(
                  stmt.node, stmt.attr_key, stmt.value, body
              )
          return None
  
      def _ftransform(f, mod, ctx):
          return f.with_body(
              tvm.tir.stmt_functor.ir_transform(
                  f.body, _do_fold, None, ["tir.AttrStmt"]
              )
          )
  
      return tvm.tir.transform.prim_func_pass(
          _ftransform, opt_level=0, name="tir.vta.FoldUopLoop"
      )
  
file_path: tvm/vta/python/vta/transform.py
imports: import tvm
  from tvm import te
  from tvm.topi import utils
  from .environment import get_env
label_desc: Use image node_id 2 for calling a global function with example usage: FoldUopLoop() and returns: tvm, None, f, None,  body,  None, op, tvm, tvm, list, stmt, op, op