code: def decl_winograd_gemm(cfg, A, B):
      """Declare a winograd GEMM for Mali Bifrost GPUs
      Winograd uses batched GEMM, however the input tensors are 4D
      This wraps decl_batched_gemm to provide it with 3D tensors
  
      Parameters
      ----------
      cfg : Config
          Schedule configuration
  
      A : tvm.te.Tensor
          4D Tensor, shape [a, a, n, k]
  
      B : tvm.te.Tensor
          4D Tensor, shape [a * a, k, m]
  
      Returns
      -------
  
      """
      alpha = utils.get_const_int(A.shape[0])
      n = utils.get_const_int(A.shape[2])
      k = utils.get_const_int(A.shape[3])
  
      A_3D = te.compute(
          (alpha * alpha, n, k),
          lambda b, i, j: A[b // alpha][b % alpha][i][j],
          name="A_3D",
      )
  
      C = decl_batched_gemm(cfg, A_3D, B)
      return A_3D, C
  
file_path: tvm/python/tvm/topi/bifrost/gemm.py
imports: from tvm import te
  from .transforms import tile_and_bind, tile_and_bind3d, interleave_transpose, transpose_interleave
  from ..None import utils
label_desc: Use image node_id 3 for calling a global function with example usage: decl_winograd_gemm(cfg, A, B) and returns: A_3D,  C