code: def _assign_poolinfos_to_allocates_in_primfunc(primfunc, pool_infos):
      """helper to assing poolinfos to allocate nodes in a tir.PrimFunc"""
  
      def set_poolinfos(stmt):
          if isinstance(stmt, tvm.tir.Allocate):
              return tvm.tir.Allocate(
                  buffer_var=stmt.buffer_var,
                  dtype=stmt.dtype,
                  extents=stmt.extents,
                  condition=stmt.condition,
                  body=stmt.body,
                  annotations={
                      tvm.tir.usmp.utils.CANDIDATE_MEMORY_POOL_ATTR: pool_infos
                  },
              )
  
      return primfunc.with_body(
          stmt_functor.ir_transform(primfunc.body, None, set_poolinfos)
      )
  
imports: import pytest
  import tvm
  from tvm import tir, script
  from tvm.script import tir
  from tvm.tir import stmt_functor
  from tvm.tir.usmp import utils
  from tvm.target import Target
  from tvm import WorkspacePoolInfo, PoolInfoProperties
label_desc: Use image node_id 4 for calling a global function with example usage: _assign_poolinfos_to_allocates_in_primfunc(primfunc, pool_infos) and returns: primfunc, tvm