code: def matmul_4d_prune_config(
      configs, named_args, skip_autotune=SKIP_AUTOTUNE
  ):
      if skip_autotune:
          configs = [configs[0]]
      else:
          device = torch.cuda.current_device()  # ignore-cuda
          capability = torch.cuda.get_device_capability()  # ignore-cuda
          # BLOCK_M, BLOCK_N, BLOCK_K, SPLIT_K, num_warps, num_stages
          dtsize = named_args["a_ptr"].element_size()
          dtype = named_args["a_ptr"].dtype
  
          # make sure we have enough smem
          pruned_configs = []
          for config in configs:
              kw = config.kwargs
              BLOCK_M, BLOCK_N, BLOCK_K, num_stages = (
                  kw["BLOCK_SIZE_M"],
                  kw["BLOCK_SIZE_N"],
                  kw["BLOCK_SIZE_K"],
                  config.num_stages,
              )
  
              max_shared_memory = (
                  triton.runtime.driver.utils.get_device_properties(
                      device
                  )["max_shared_mem"]
              )
              required_shared_memory = (
                  (BLOCK_M + BLOCK_N) * BLOCK_K * num_stages * dtsize
              )
              if required_shared_memory <= max_shared_memory:
                  pruned_configs.append(config)
          configs = pruned_configs
      return configs
  
file_path: DeepSpeed/deepspeed/ops/transformer/inference/triton/triton_matmul_kernel.py
imports: import triton
  import triton.language
  from .gelu import gelu_functor
  import torch
label_desc: Use image node_id 4 for calling a global function with example usage: matmul_4d_prune_config(configs, named_args, skip_autotune) and returns: configs