code: def gemm_acc_nx16_int8_int8_int32(dtype, rows):
      """
      Int8 nx16 matrix multiplication and accumulation using sdot/udot instructions
      This function takes two arrays of int8 datatype -- A[n][4] and
      B[4][16] and produces a rowsx16 matrix which is equal to A*B'
      The pseudo code is as follows.
  
      .. code-block:: c
  
          void mmla_nx16_int8_int8_int32(int8 A[n][16], int8 B[4][16][4], int32 output[n][16]){
              for (int i = 0; i < n; i++){
                  for (int j = 0; j < 16; j++){
                      for (int k = 0; k < 16; k++){
                          out[i][j] += A[i][k] * B[k//4][j][k%4]
                      }
                  }
              }
          }
  
      Notes:
          * The tile size of B is 16x4. Since the reduction variable k moves between 0 and 16
            we need 4 tiles of B to compute a single row of the output. The first 4 values of
            k will be fetched from B[0][j][k], the second batch of 4 from B[1][j][k] and so on
          * The tiling strategy is picked to maximize register usage.
  
      Parameters
      ----------
      dtype : str, {"uint8", "int8"}
          Whether it works on unsigned int or signed int
      rows : int
          Number of the output rows "n"
  
      Returns
      -------
      intrin : TensorIntrin
          The Arm TensorIntrin that can be used in tensorizing schedule
      """
      assert dtype in ["uint8", "int8"]
      A = te.placeholder((rows, 16), dtype, name="A")
      B = te.placeholder((4, 16, 4), dtype, name="B")
      dtype_vec = dtype + "x16"
      idxm = tvm.tir.indexmod
      k = te.reduce_axis((0, 16), name="k")
      C = te.compute(
          (rows, 16),
          lambda i, j: te.sum(
              A[i, k].astype("int32")
              * B[k // 4, j, idxm(k, 4)].astype("int32"),
              axis=k,
          ),
          name="C",
      )
  
      aa_buffer = tvm.tir.decl_buffer(
          A.shape,
          dtype,
          name="aa_buffer",
          offset_factor=1,
          strides=[te.var("sa"), 1],
      )
      bb_buffer = tvm.tir.decl_buffer(
          B.shape,
          dtype,
          name="bb_buffer",
          offset_factor=1,
          strides=[te.var("sb0"), te.var("sb1"), 1],
      )
      cc_buffer = tvm.tir.decl_buffer(
          C.shape,
          dtype="int32",
          name="cc_buffer",
          offset_factor=1,
          strides=[te.var("sc"), 1],
      )
  
      llvm_intrin = (
          "llvm.aarch64.neon.sdot"
          if dtype == "int8"
          else "llvm.aarch64.neon.udot"
      )
  
      def _intrin_func(ins, outs):
          def _instr(index):
              ib = tvm.tir.ir_builder.create()
              if index == 1:
                  for i in range(0, rows):
                      ib.emit(
                          outs[0].vstore(
                              [i, 0], tvm.tir.const(0, "int32x16")
                          )
                      )
                  return ib.get()
              # Iterate on the number of rows of the output
              for k in range(0, rows):
                  # Load 16 elements of A
                  # vec_a = [a, b, c, d, e, f, g, h, l, m, n, o, p, q, r, s];
                  vec_a = ins[0].vload([k, 0], dtype_vec)
  
                  # Iterate over each of the 4 rowsx4 tiles of the output
                  for j in range(0, 4):
                      # Accumulate over each of the 4 (16x4) tiles contained in B
                      for i in range(0, 4):
                          # Replicate a single 4-element group of A (A[k, i:i+4])
                          vec_aa = select_word(vec_a, i, dtype_vec)
  
                          # Load 4 rows (each rows with 4 elements) from B (B[i:i+4, j:j+4])
                          # vec_b = [0, 16, 32, 48,
                          #          1, 17, 33, 49,
                          #          2, 18, 34, 50,
                          #          3, 19, 35, 51,];
                          vec_b = ins[1].vload([i, 4 * j, 0], dtype_vec)
  
                          # Accumulate in the correct part of the output
                          vec_c = outs[0].vload([k, 4 * j], "int32x4")
  
                          # Compute the dot product between the rowsx4 tile
                          # from A and the 4x4 tile from B
                          #
                          # For instance, for i=0, we have:
                          # sdot(vec_aa[0], vec_b) = [a*0+b*16+c*32+d*48,
                          #                           a*1+b*17+c*33+d*49,
                          #                           a*2+b*18+c*34+d*50,
                          #                           a*3+b*19+c*35+d*51]
                          vdot = tvm.tir.call_llvm_intrin(
                              "int32x4",
                              llvm_intrin,
                              tvm.tir.const(3, "uint32"),
                              vec_c,
                              vec_b,
                              vec_aa,
                          )
                          ib.emit(outs[0].vstore([k, 4 * j], vdot))
              return ib.get()
  
          # body, reset, update
          return _instr(0), _instr(1), _instr(2)
  
      buffer_params = {"offset_factor": 1}
      return te.decl_tensor_intrin(
          C.op,
          _intrin_func,
          binds={A: aa_buffer, B: bb_buffer, C: cc_buffer},
          default_buffer_params=buffer_params,
      )
  
file_path: tvm/python/tvm/topi/arm_cpu/tensor_intrin.py
imports: import tvm
  from tvm import te
  from tvm.ir import register_intrin_lowering
label_desc: Use image node_id 6 for calling a global function with example usage: gemm_acc_nx16_int8_int8_int32(dtype, rows) and returns: te, _instr,  _instr,  _instr, ib, ib