code: def decl_gemm(cfg, A, B):
      """Declare a single GEMM computation for Mali Bifrost GPUs
  
      Parameters
      ----------
      cfg : Config
          Schedule configuration
  
      A : tvm.te.Tensor
          2D Tensor, shape [n, k]
  
      B : tvm.te.Tensor
          2D Tensor, shape [k, m]
  
      Returns
      -------
      C : tvm.te.Tensor
          2D Tensor, shape [n, m]
      """
  
      cfg.define_knob(
          "work_group_x", [1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 64]
      )
      cfg.define_knob(
          "work_group_y", [1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 64]
      )
      cfg.define_knob("unroll_k_factor", [1, 2, 4])
      cfg.define_knob("A_interleave", [1, 4, 8, 16, 24, 32, 48, 64])
      cfg.define_knob("B_interleave", [1, 4, 8, 16, 32])
      cfg.define_knob("split_k_factor", [1, 4, 16])
  
      # Mutual k axis must be of equal extent
      assert utils.get_const_int(A.shape[1]) == utils.get_const_int(
          B.shape[0]
      )
      n = A.shape[0]
      m = B.shape[1]
      k_size = utils.get_const_int(A.shape[1])
      unroll_gemm = cfg["split_k_factor"].val
      if unroll_gemm == 1:
          # No unrolling case must have the same set of tensors to keep scheduling consistent
          # Create identity tensors to take the place of A_unrolled, B_unrolled and R
          A_unrolled = te.compute(
              (n, k_size), lambda i, j: A[i, j], name="A_unrolled"
          )
          B_unrolled = te.compute(
              (k_size, m), lambda i, j: B[i, j], name="B_unrolled"
          )
  
          # Declare standard GEMM
          k = te.reduce_axis((0, A.shape[1]), name="k")
          C = te.compute(
              (n, m),
              lambda i, j: te.sum(
                  A_unrolled[i, k] * B_unrolled[k, j], axis=k
              ),
              name="C",
          )
  
          R = te.compute((n, m), lambda i, j: C[i, j], name="R")
  
      else:
          unrolled_k_size = k_size // unroll_gemm
  
          # Unroll the two input matrices along the shared k axis
          A_unrolled = te.compute(
              (unroll_gemm, n, unrolled_k_size),
              lambda b, i, j: A[i][unrolled_k_size * b + j],
              name="A_unrolled",
          )
  
          B_unrolled = te.compute(
              (unroll_gemm, unrolled_k_size, m),
              lambda b, i, j: B[unrolled_k_size * b + i][j],
              name="B_unrolled",
          )
  
          # Declare a batched GEMM
          k = te.reduce_axis((0, unrolled_k_size), name="k")
          C = te.compute(
              (unroll_gemm, n, m),
              lambda b, i, j: te.sum(
                  A_unrolled[b][i][k] * B_unrolled[b][k][j], axis=k
              ),
              name="C",
          )
  
          # Then declare a reduction to reduce the sub matrices
          k = te.reduce_axis((0, unroll_gemm), name="k")
          R = te.compute(
              (n, m), lambda i, j: te.sum(C[k][i][j], axis=k), name="R"
          )
  
      return R
  
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 1 for calling a global function with example usage: decl_gemm(cfg, A, B) and returns: R