code: def assign_poolinfos_to_allocates_in_irmodule(mod, pool_infos):
      """Helper to assign poolinfos to allocate nodes in a IRModule"""
      ret = tvm.IRModule()
      for global_var, basefunc in mod.functions.items():
          if isinstance(basefunc, tvm.tir.PrimFunc):
              ret[
                  global_var
              ] = assign_poolinfos_to_allocates_in_primfunc(
                  basefunc, pool_infos
              )
      return ret
  
imports: import sys
  import pytest
  import tvm
  from tvm import PoolInfoProperties, WorkspacePoolInfo
  from tvm.script import tir, ir
  from tvm.target import Target
  from tvm.tir import stmt_functor
  from tvm.tir.usmp import utils
label_desc: Use image node_id 3 for calling a global function with example usage: assign_poolinfos_to_allocates_in_irmodule(mod, pool_infos) and returns: ret