code: def _batch_matmul_legalize(attrs, inputs, arg_types):
      """Legalizes batch_matmul op.
  
      Parameters
      ----------
      attrs : tvm.ir.Attrs
          Attributes of current convolution
      inputs : list of tvm.relay.Expr
          The args of the Relay expr to be legalized
      arg_types : list of types
          List of input and output types
  
      Returns
      -------
      result : tvm.relay.Expr
          The legalized expr
      """
      # Collect the input tensors.
      x_tensor, y_tensor = arg_types[0], arg_types[1]
      dtype = x_tensor.dtype
  
      if attrs.transpose_a:
          B, K, M = x_tensor.shape
      else:
          B, M, K = x_tensor.shape
  
      if attrs.transpose_b:
          B, N, K = y_tensor.shape
      else:
          B, K, N = y_tensor.shape
  
      # Collect the output tensor.
      output_tensor = arg_types[2]
  
      # Collect the input exprs.
      x, y = inputs
  
      if (
          isinstance(B, tir.expr.Any)
          or isinstance(M, tir.expr.Any)
          or isinstance(K, tir.expr.Any)
          or isinstance(N, tir.expr.Any)
      ):
          # Dynamic shape do not support alter op layout now
          return None
  
      M = M.value
      K = K.value
      N = N.value
  
      # Pad input and output channels to use tensorcore schedule.
      if dtype in ["float16", "int8", "uint8"]:
          # The shape of (M, K, N) must be multiple of (16, 16, 16) or (32, 16, 8) or (8, 16, 32)
          if (
              (M % 8 == 0 and K % 16 == 0 and N % 32 == 0)
              or (M % 16 == 0 and K % 16 == 0 and N % 16 == 0)
              or (M % 32 == 0 and K % 16 == 0 and N % 8 == 0)
          ):
              # no need to pad
              return None
          candidates = [(16, 16, 16), (32, 16, 8), (8, 16, 32)]
      elif dtype in ["int4", "uint4"]:
          if M % 8 == 0 and K % 32 == 0 and N % 8 == 0:
              # no need to pad
              return None
  
          candidates = [(8, 32, 8)]
      else:
          return None
  
      (dm, dk, dn), extra_flops = pad_to_tensorcore(M, K, N, candidates)
  
      if extra_flops > 2:
          logger.info(
              "batch_matmul pad_to_tensorcore skipped, extra_flops %s",
              extra_flops,
          )
          return None
  
      logger.info(
          "batch_matmul pad_to_tensorcore, extra_flops %s", extra_flops
      )
  
      if attrs.transpose_a:
          pad_width = ((0, 0), (0, dk), (0, dm))
      else:
          pad_width = ((0, 0), (0, dm), (0, dk))
  
      x_ = relay.nn.pad(x, pad_width=pad_width) if dm or dk else x
  
      if attrs.transpose_b:
          pad_width = ((0, 0), (0, dn), (0, dk))
      else:
          pad_width = ((0, 0), (0, dk), (0, dn))
  
      y_ = relay.nn.pad(y, pad_width=pad_width) if dn or dk else y
  
      out_ = relay.nn.batch_matmul(x_, y_, **attrs)
  
      out = (
          relay.strided_slice(
              out_,
              begin=[0, 0, 0],
              end=[x.value for x in output_tensor.shape],
          )
          if dm or dn
          else out_
      )
      return out
  
file_path: tvm/python/tvm/topi/cuda/tensorcore_alter_op.py
imports: import logging
  import math
  from tvm import relay, tir
  from ..None import nn
label_desc: Use image node_id 1 for calling a global function with example usage: _batch_matmul_legalize(attrs, inputs, arg_types) and returns: out, None, None, None, None, None