code: def test_tensorize_matmul():
      n = 1024
      m = n
      l = n
      A = te.placeholder((n, l), name="A")
      B = te.placeholder((m, l), name="B")
      k = te.reduce_axis((0, l), name="k")
      C = te.compute(
          (n, m),
          lambda i, j: te.sum(B[j, k] * A[i, k], axis=k),
          name="C",
      )
  
      def check(factor):
          s = te.create_schedule(C.op)
          x, y = C.op.axis
          yo, yi = s[C].split(y, factor=factor)
          gemv = intrin_gemv(factor, l)
          s[C].tensorize(yi, gemv)
          s = s.normalize()
          dom_map = tvm.te.schedule.InferBound(s)
          finfer = tvm.get_global_func("test.op.InferTensorizeRegion")
          out_dom, in_dom = finfer(s[C], dom_map)
          assert tvm.ir.structural_equal(out_dom[x].extent, 1)
          assert tvm.ir.structural_equal(out_dom[y].extent, factor)
          assert tvm.ir.structural_equal(out_dom[y].min, yo * factor)
          fmatch = tvm.get_global_func("test.op.MatchTensorizeBody")
          body = fmatch(s[C], out_dom, in_dom, gemv)
          ana = tvm.arith.Analyzer()
  
          assert tvm.ir.structural_equal(
              ana.simplify(body[0]), ana.simplify(gemv.op.body[0])
          )
          stmt = tvm.te.schedule.ScheduleOps(s, dom_map)
          tvm.lower(s, [A, B, C])
  
      def check_rfactor(factor, rfactor):
          s = te.create_schedule(C.op)
          x, y = C.op.axis
          rk = C.op.reduce_axis[0]
          yo, yi = s[C].split(y, factor=factor)
          ro, ri = s[C].split(rk, factor=rfactor)
          s[C].reorder(yo, ro, yi, ri)
          gemv = intrin_gemv(factor, rfactor)
          s[C].tensorize(yi, gemv)
          s = s.normalize()
          dom_map = tvm.te.schedule.InferBound(s)
          finfer = tvm.get_global_func("test.op.InferTensorizeRegion")
          out_dom, in_dom = finfer(s[C], dom_map)
          assert tvm.ir.structural_equal(out_dom[x].extent, 1)
          assert tvm.ir.structural_equal(out_dom[y].extent, factor)
          assert tvm.ir.structural_equal(out_dom[y].min, yo * factor)
          fmatch = tvm.get_global_func("test.op.MatchTensorizeBody")
          body = fmatch(s[C], out_dom, in_dom, gemv)
          ana = tvm.arith.Analyzer()
          assert tvm.ir.structural_equal(
              ana.simplify(body[0]), ana.simplify(gemv.op.body[0])
          )
          stmt = tvm.te.schedule.ScheduleOps(s, dom_map)
          tvm.lower(s, [A, B, C])
  
      def check_rfactor_no_reset(factor, rfactor):
          s = te.create_schedule(C.op)
          x, y = C.op.axis
          rk = C.op.reduce_axis[0]
          yo, yi = s[C].split(y, factor=factor)
          ro, ri = s[C].split(rk, factor=rfactor)
          s[C].reorder(yo, ro, yi, ri)
          gemv = intrin_gemv_no_reset(factor, rfactor)
          s[C].tensorize(yi, gemv)
          s = s.normalize()
          dom_map = tvm.te.schedule.InferBound(s)
          finfer = tvm.get_global_func("test.op.InferTensorizeRegion")
          out_dom, in_dom = finfer(s[C], dom_map)
          assert tvm.ir.structural_equal(out_dom[x].extent, 1)
          assert tvm.ir.structural_equal(out_dom[y].extent, factor)
          assert tvm.ir.structural_equal(out_dom[y].min, yo * factor)
          fmatch = tvm.get_global_func("test.op.MatchTensorizeBody")
          body = fmatch(s[C], out_dom, in_dom, gemv)
          ana = tvm.arith.Analyzer()
          assert tvm.ir.structural_equal(
              ana.simplify(body[0]), ana.simplify(gemv.op.body[0])
          )
          stmt = tvm.te.schedule.ScheduleOps(s, dom_map)
          tvm.lower(s, [A, B, C])
  
      def check_rfactor_no_reset_multi_reduction(factor, rfactor):
          s = te.create_schedule(C.op)
          x, y = C.op.axis
          rk = C.op.reduce_axis[0]
          yo, yi = s[C].split(y, factor=factor)
          ro, ri = s[C].split(rk, factor=rfactor)
          roo, roi = s[C].split(ro, factor=2)
          s[C].reorder(yo, roo, roi, yi, ri)
          gemv = intrin_gemv_no_reset(factor, rfactor)
          s[C].tensorize(yi, gemv)
          s = s.normalize()
          dom_map = tvm.te.schedule.InferBound(s)
          finfer = tvm.get_global_func("test.op.InferTensorizeRegion")
          out_dom, in_dom = finfer(s[C], dom_map)
          assert tvm.ir.structural_equal(out_dom[x].extent, 1)
          assert tvm.ir.structural_equal(out_dom[y].extent, factor)
          assert tvm.ir.structural_equal(out_dom[y].min, yo * factor)
          fmatch = tvm.get_global_func("test.op.MatchTensorizeBody")
          body = fmatch(s[C], out_dom, in_dom, gemv)
          ana = tvm.arith.Analyzer()
          assert tvm.ir.structural_equal(
              ana.simplify(body[0]), ana.simplify(gemv.op.body[0])
          )
          stmt = tvm.te.schedule.ScheduleOps(s, dom_map)
          tvm.lower(s, [A, B, C])
  
      check(16)
      check_rfactor(16, 16)
      check_rfactor_no_reset(16, 16)
      check_rfactor_no_reset_multi_reduction(16, 16)
  
file_path: tvm/tests/python/te/test_te_schedule_tensorize.py
imports: import tvm
  from tvm import te
label_desc: Use image node_id 5 for calling a global function with example usage: test_tensorize_matmul() without return types