code: def test_task_extraction_cpu(params):
      ishape = (1, 3, 224, 224)
      w1shape = (32, 3, 3, 3)
      w2shape = (32, 32, 3, 3)
      dtype = "float32"
      target = tvm.target.Target("llvm")
  
      def get_func():
          data = relay.var("data", shape=(ishape), dtype=dtype)
          weight1 = relay.var("weight1", shape=(w1shape), dtype=dtype)
          weight2 = relay.var("weight2", shape=(w2shape), dtype=dtype)
  
          conv2d = relay.nn.conv2d(
              data, weight1, kernel_size=(3, 3), padding=(1, 1)
          )
          relu = relay.nn.relu(conv2d)
          conv2d = relay.nn.conv2d(
              relu, weight2, kernel_size=(3, 3), padding=(1, 1)
          )
          out = relay.nn.relu(conv2d)
          return relay.Function([data, weight1, weight2], out)
  
      def get_fused_func():
          data = relay.var("data", shape=(ishape), dtype=dtype)
          weight1 = relay.var("weight1", shape=(w1shape), dtype=dtype)
          weight2 = relay.var("weight2", shape=(w2shape), dtype=dtype)
  
          fused_func = get_func()
  
          # Set to primitive to keep fuse_ops untouch.
          fused_func = fused_func.with_attr(
              "Primitive", tvm.tir.IntImm("int32", 1)
          )
  
          call = relay.Call(fused_func, [data, weight1, weight2])
          return relay.Function([data, weight1, weight2], call)
  
      def get_simple_func():
          data = relay.var(
              "data", relay.TensorType((1, 2, 3), "float32")
          )
          out = relay.image.affine_grid(data, (150, 150))
          return relay.Function([data], out)
  
      def get_shape_of_func():
          data = relay.var(
              "data", shape=(relay.Any(), 28, 28), dtype="float32"
          )
          out = relay.shape_of(data)
          return relay.Function([data], out)
  
      def get_func_with_dynamic_shape():
          data = relay.var(
              "data", shape=(relay.Any(), 32), dtype="float32"
          )
          out = relay.max(data)
          return relay.Function(relay.analysis.free_vars(out), out)
  
      def get_func_with_control_flow():
          data = relay.var("data", shape=(1, 3, 224, 224))
          weight = relay.var("weight", shape=(3, 3, 3, 3))
          eq1 = relay.var("e1", shape=[], dtype="float32")
          eq2 = relay.var("e2", shape=[], dtype="float32")
          eq = relay.equal(eq1, eq2)
  
          true_branch = relay.zeros(
              shape=(1, 3, 224, 224), dtype="float32"
          )
          false_branch = relay.nn.conv2d(
              data,
              weight,
              kernel_size=(3, 3),
              channels=3,
              padding=(1, 1),
          )
          false_branch = relay.nn.conv2d(
              false_branch,
              weight,
              kernel_size=(3, 3),
              channels=3,
              padding=(1, 1),
          )
          ife = relay.If(eq, true_branch, false_branch)
          out = relay.erf(ife)
          return relay.Function([data, weight, eq1, eq2], out)
  
      def get_func_with_unsupported_op():
          def get_postproc_func():
              data = relay.var("data", shape=((1, 3, 6)), dtype=dtype)
              out = relay.nn.relu(data)
              func = relay.Function([data], out)
              func = func.with_attr(
                  "Primitive", tvm.tir.IntImm("int32", 1)
              )
              return func
  
          cls_prob = relay.var(
              "cls_prob", relay.ty.TensorType((1, 3, 3), "float32")
          )
          loc_pred = relay.var(
              "loc_pred", relay.ty.TensorType((1, 3 * 4), "float32")
          )
          anchors = relay.var(
              "anchors", relay.ty.TensorType((1, 3, 4), "float32")
          )
  
          mtl = relay.vision.multibox_transform_loc(
              cls_prob=cls_prob, loc_pred=loc_pred, anchor=anchors
          )
          nms = relay.vision.non_max_suppression(
              mtl[0], mtl[1], mtl[0], return_indices=False
          )
          out = relay.Call(get_postproc_func(), [nms])
          return relay.Function([cls_prob, loc_pred, anchors], out)
  
      func_map = {
          "basic_func": get_func,
          "fused_func": get_fused_func,
          "simple_func": get_simple_func,
          "shape_of_func": get_shape_of_func,
          "dyn_shape_func": get_func_with_dynamic_shape,
          "control_flow_func": get_func_with_control_flow,
          "func_w_unsupported_op": get_func_with_unsupported_op,
      }
  
      def verify_task_extraction(
          func_name, expected_task, include_simple_tasks=False
      ):
          func = func_map[func_name]()
          mod = tvm.IRModule.from_expr(func)
          tasks, task_weights = auto_scheduler.extract_tasks(
              mod["main"],
              None,
              target,
              include_simple_tasks=include_simple_tasks,
          )
  
          assert len(tasks) == expected_task
          assert len(task_weights) == expected_task
  
      verify_task_extraction(*params)
  
file_path: tvm/tests/python/relay/test_auto_scheduler_task_extraction.py
imports: import json
  import tempfile
  import pytest
  import tvm.relay.testing
  import tvm.testing
  from tvm import _ffi
  from tvm import auto_scheduler, relay
label_desc: Use image node_id 3 for calling a global function with example usage: test_task_extraction_cpu(params) and returns: relay, relay, relay, relay, relay, relay, relay, func