zoukankan      html  css  js  c++  java
  • CPU的自动调度矩阵乘法

    CPU的自动调度矩阵乘法

    这是一个有关如何对CPU使用自动调度程序的文档。

    与依靠手动模板定义搜索空间的基于模板的autotvm不同,自动调度程序不需要任何模板。用户只需要编写计算声明,而无需任何调度命令或模板。自动调度程序可以自动生成较大的搜索空间,并在该空间中找到良好的调度。

    本文以矩阵乘法为例。

    注意,本文无法在Windows或最新版本的macOS上运行。要使其运行,需要将本文的内容包装在一个if __name__ == "__main__":块中。

    import os

     

    import numpy as np

    import tvm

    from tvm import te, auto_scheduler

    定义计算

    首先,定义带有偏差加法的矩阵的计算。该函数应返回输入/输出张量的列表。通过这些张量,自动调度器可以获取整个计算图。

    @auto_scheduler.register_workload

    def matmul_add(N, L, M, dtype):

        A = te.placeholder((N, L), name="A", dtype=dtype)

        B = te.placeholder((L, M), name="B", dtype=dtype)

        C = te.placeholder((N, M), name="C", dtype=dtype)

     

        k = te.reduce_axis((0, L), name="k")

        matmul = te.compute(

            (N, M),

            lambda i, j: te.sum(A[i, k] * B[k, j], axis=k),

            name="matmul",

            attrs={"layout_free_placeholders": [B]},  # enable automatic layout transform for tensor B

        )

        out = te.compute((N, M), lambda i, j: matmul[i, j] + C[i, j], name="out")

     

        return [A, B, C, out]

    创建搜索任务

    然后,创建一个搜索任务,其中N = L = M = 1024且dtype =“ float32”。如果计算机支持avx指令,可以

    • 将下面的“ llvm”替换为“ llvm -mcpu = core-avx2”以启用AVX2
    • 将下面的“ llvm”替换为“ llvm -mcpu = skylake-avx512”以启用AVX-512

    target = tvm.target.Target("llvm")

    N = L = M = 1024

    task = tvm.auto_scheduler.SearchTask(func=matmul_add, args=(N, L, M, "float32"), target=target)

     

    # Inspect the computational graph

    print("Computational DAG:")

    print(task.compute_dag)

    出:

    Computational DAG:

    A = PLACEHOLDER [1024, 1024]

    B = PLACEHOLDER [1024, 1024]

    matmul(i, j) += (A[i, k]*B[k, j])

    C = PLACEHOLDER [1024, 1024]

    out(i, j) = (matmul[i, j] + C[i, j])

    接下来,为自动调度程序设置参数。

    • num_measure_trials是在搜索过程中可以使用的测量试验的数量。为了快速演示,在本文中仅进行10次试用。实际上,1000是搜索收敛的一个很好的值。可以根据自己的时间预算进行更多试验。
    • 此外,还用RecordToFile将测量记录转储到文件matmul.json。测量记录可用于最好地查询历史记录,恢复搜索以及以后进行更多分析。
    • 查看更多参数auto_scheduler.TuningOptions

    log_file = "matmul.json"

    tune_option = auto_scheduler.TuningOptions(

        num_measure_trials=10,

        measure_callbacks=[auto_scheduler.RecordToFile(log_file)],

        verbose=2,

    )

    运行搜索

    现在准备好所有输入。开始搜索,让自动调度程序发挥作用。经过一些测量试验后,可以从日志文件中加载最佳调度并应用它。

    # Run auto-tuning (search)

    task.tune(tune_option)

    # Apply the best schedule

    sch, args = task.apply_best(log_file)

    出:

    *T*T*T*T*T*T*T*T*T*T

    可以降低调度,以便在自动调度后查看IR。自动调度程序正确执行优化,包括多层平铺,布局转换,并行化,矢量化,展开和运算符融合。

    print("Lowered TIR:")

    print(tvm.lower(sch, args, simple_mode=True))

    输出:

    Lowered TIR:

    primfn(A_1: handle, B_1: handle, C_1: handle, out_1: handle) -> ()

      attr = {"global_symbol": "main", "tir.noalias": True}

      buffers = {out: Buffer(out_2: Pointer(float32), float32, [1024, 1024], []),

                 C: Buffer(C_2: Pointer(float32), float32, [1024, 1024], []),

                 B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], []),

                 A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], [])}

      buffer_map = {A_1: A, B_1: B, C_1: C, out_1: out} {

      attr [auto_scheduler_layout_transform: Pointer(float32)] "storage_scope" = "global";

      allocate(auto_scheduler_layout_transform, float32, [1048576]) {

        for (ax0.ax1.fused.ax2.fused.ax3.fused.ax4.fused.ax5.fused.ax6.fused: int32, 0, 131072) "parallel" {

          for (ax7: int32, 0, 8) {

            auto_scheduler_layout_transform[((ax0.ax1.fused.ax2.fused.ax3.fused.ax4.fused.ax5.fused.ax6.fused*8) + ax7)] = (float32*)B_2[(((floormod(ax0.ax1.fused.ax2.fused.ax3.fused.ax4.fused.ax5.fused.ax6.fused, 1024)*1024) + (floordiv(ax0.ax1.fused.ax2.fused.ax3.fused.ax4.fused.ax5.fused.ax6.fused, 1024)*8)) + ax7)]

          }

        }

        for (i.outer.outer.j.outer.outer.fused: int32, 0, 16384) "parallel" {

          attr [matmul: Pointer(float32)] "storage_scope" = "global";

          allocate(matmul, float32x8, [4]);

          for (i.outer.inner: int32, 0, 2) {

            matmul[ramp(0, 1, 8)] = broadcast(0f32, 8)

            matmul[ramp(8, 1, 8)] = broadcast(0f32, 8)

            matmul[ramp(16, 1, 8)] = broadcast(0f32, 8)

            matmul[ramp(24, 1, 8)] = broadcast(0f32, 8)

            for (k.outer: int32, 0, 256) {

              for (k.inner: int32, 0, 4) {

                matmul[ramp(0, 1, 8)] = ((float32x8*)matmul[ramp(0, 1, 8)] + (broadcast((float32*)A_2[((((floordiv(i.outer.outer.j.outer.outer.fused, 128)*8192) + (i.outer.inner*4096)) + (k.outer*4)) + k.inner)], 8)*(float32x8*)auto_scheduler_layout_transform[ramp((((floormod(i.outer.outer.j.outer.outer.fused, 128)*8192) + (k.outer*32)) + (k.inner*8)), 1, 8)]))

                matmul[ramp(8, 1, 8)] = ((float32x8*)matmul[ramp(8, 1, 8)] + (broadcast((float32*)A_2[(((((floordiv(i.outer.outer.j.outer.outer.fused, 128)*8192) + (i.outer.inner*4096)) + (k.outer*4)) + k.inner) + 1024)], 8)*(float32x8*)auto_scheduler_layout_transform[ramp((((floormod(i.outer.outer.j.outer.outer.fused, 128)*8192) + (k.outer*32)) + (k.inner*8)), 1, 8)]))

                matmul[ramp(16, 1, 8)] = ((float32x8*)matmul[ramp(16, 1, 8)] + (broadcast((float32*)A_2[(((((floordiv(i.outer.outer.j.outer.outer.fused, 128)*8192) + (i.outer.inner*4096)) + (k.outer*4)) + k.inner) + 2048)], 8)*(float32x8*)auto_scheduler_layout_transform[ramp((((floormod(i.outer.outer.j.outer.outer.fused, 128)*8192) + (k.outer*32)) + (k.inner*8)), 1, 8)]))

                matmul[ramp(24, 1, 8)] = ((float32x8*)matmul[ramp(24, 1, 8)] + (broadcast((float32*)A_2[(((((floordiv(i.outer.outer.j.outer.outer.fused, 128)*8192) + (i.outer.inner*4096)) + (k.outer*4)) + k.inner) + 3072)], 8)*(float32x8*)auto_scheduler_layout_transform[ramp((((floormod(i.outer.outer.j.outer.outer.fused, 128)*8192) + (k.outer*32)) + (k.inner*8)), 1, 8)]))

              }

            }

            for (i.inner: int32, 0, 4) {

              out_2[ramp(((((floordiv(i.outer.outer.j.outer.outer.fused, 128)*8192) + (i.outer.inner*4096)) + (i.inner*1024)) + (floormod(i.outer.outer.j.outer.outer.fused, 128)*8)), 1, 8)] = ((float32x8*)matmul[ramp((i.inner*8), 1, 8)] + (float32x8*)C_2[ramp(((((floordiv(i.outer.outer.j.outer.outer.fused, 128)*8192) + (i.outer.inner*4096)) + (i.inner*1024)) + (floormod(i.outer.outer.j.outer.outer.fused, 128)*8)), 1, 8)])

            }

          }

        }

      }

    }

    检查正确性并评估性能

    构建二进制文件并检查其正确性和性能。

    func = tvm.build(sch, args, target)

    a_np = np.random.uniform(size=(N, L)).astype(np.float32)

    b_np = np.random.uniform(size=(L, M)).astype(np.float32)

    c_np = np.random.uniform(size=(N, M)).astype(np.float32)

    out_np = a_np.dot(b_np) + c_np

     

    ctx = tvm.cpu()

    a_tvm = tvm.nd.array(a_np, ctx=ctx)

    b_tvm = tvm.nd.array(b_np, ctx=ctx)

    c_tvm = tvm.nd.array(c_np, ctx=ctx)

    out_tvm = tvm.nd.empty(out_np.shape, ctx=ctx)

    func(a_tvm, b_tvm, c_tvm, out_tvm)

     

    # Check results

    np.testing.assert_allclose(out_np, out_tvm.asnumpy(), rtol=1e-3)

     

    # Evaluate execution time.

    evaluator = func.time_evaluator(func.entry_name, ctx, min_repeat_ms=500)

    print(

        "Execution time of this operator: %.3f ms"

        % (np.median(evaluator(a_tvm, b_tvm, c_tvm, out_tvm).results) * 1000)

    )

    出:

    Execution time of this operator: 22.426 ms

    使用记录文件

    搜索期间,所有测量记录都将转储到记录文件“ matmul.json”中。测量记录可用于重新应用搜索结果,继续搜索以及执行其它分析。

    这是一个示例,其中从文件加载最佳调度,并打印等效的python调度API。这可用于调试和学习自动调度程序的行为。

    print("Equivalent python schedule:")

    print(task.print_best(log_file))

    出:

    Equivalent python schedule:

    matmul_i, matmul_j, matmul_k = tuple(matmul.op.axis) + tuple(matmul.op.reduce_axis)

    out_i, out_j = tuple(out.op.axis) + tuple(out.op.reduce_axis)

    matmul_i_o_i, matmul_i_i = s[matmul].split(matmul_i, factor=4)

    matmul_i_o_o_i, matmul_i_o_i = s[matmul].split(matmul_i_o_i, factor=1)

    matmul_i_o_o_o, matmul_i_o_o_i = s[matmul].split(matmul_i_o_o_i, factor=2)

    matmul_j_o_i, matmul_j_i = s[matmul].split(matmul_j, factor=8)

    matmul_j_o_o_i, matmul_j_o_i = s[matmul].split(matmul_j_o_i, factor=1)

    matmul_j_o_o_o, matmul_j_o_o_i = s[matmul].split(matmul_j_o_o_i, factor=1)

    matmul_k_o, matmul_k_i = s[matmul].split(matmul_k, factor=4)

    s[matmul].reorder(matmul_i_o_o_o, matmul_j_o_o_o, matmul_i_o_o_i, matmul_j_o_o_i, matmul_k_o, matmul_i_o_i, matmul_j_o_i, matmul_k_i, matmul_i_i, matmul_j_i)

    out_i_o_i, out_i_i = s[out].split(out_i, factor=4)

    out_i_o_o, out_i_o_i = s[out].split(out_i_o_i, factor=2)

    out_j_o_i, out_j_i = s[out].split(out_j, factor=8)

    out_j_o_o, out_j_o_i = s[out].split(out_j_o_i, factor=1)

    s[out].reorder(out_i_o_o, out_j_o_o, out_i_o_i, out_j_o_i, out_i_i, out_j_i)

    s[matmul].compute_at(s[out], out_j_o_i)

    out_i_o_o_j_o_o_fused = s[out].fuse(out_i_o_o, out_j_o_o)

    s[out].parallel(out_i_o_o_j_o_o_fused)

    s[matmul].pragma(matmul_i_o_o_o, "auto_unroll_max_step", 8)

    s[matmul].pragma(matmul_i_o_o_o, "unroll_explicit", True)

    s[matmul].vectorize(matmul_j_i)

    s[out].vectorize(out_j_i)

    一个更复杂的示例是继续搜索。在这种情况下,需要自己创建搜索策略和成本模型,并使用日志文件恢复搜索策略和成本模型的状态。在下面的示例中,恢复状态并进行5次以上的试用。

    def resume_search(task, log_file_name):

        cost_model = auto_scheduler.XGBModel()

        cost_model.update_from_file(log_file_name)

        search_policy = auto_scheduler.SketchPolicy(

            task,

            cost_model,

            init_search_callbacks=[auto_scheduler.PreloadMeasuredStates(log_file_name)],

        )

        tune_option = auto_scheduler.TuningOptions(

            num_measure_trials=5, measure_callbacks=[auto_scheduler.RecordToFile(log_file_name)]

        )

        task.tune(tune_option, search_policy=search_policy)

     

     

    # resume_search(task, log_file)

    注意

    由于python的多处理和tvm的线程池之间的冲突,因此无法在上面运行此行。运行tvm生成的二进制文件后,python的多处理库将永远挂起。必须确保在调用auot-scheduler的搜索之前,不要运行任何tvm生成的二进制文件。要运行上面的功能,应该注释掉“检查正确性和评估性能”部分中的所有代码。

    应该在应用程序中注意这个问题。对于此问题,还有其他解决方法。例如,可以启动新线程/进程(使用内置的python库线程或多线程处理),并在新线程/进程中运行tvm二进制文件。这提供了隔离,并避免了主线程/进程中的冲突。还可以将auto_scheduler.LocalRPCMeasureContext用于自动调度程序,如GPU帮助(自动调度GPU的卷积层)中所示。

    脚本的总运行时间:(1分钟50.410秒)

    https://tvm.apache.org/docs/tutorials/auto_scheduler/tune_matmul_x86.html#sphx-glr-tutorials-auto-scheduler-tune-matmul-x86-py

    人工智能芯片与自动驾驶
  • 相关阅读:
    在.NET访问MySql数据库时的几点经验(转)
    FxCop代码标准检测工具
    ASP(从前) vs ASP.NET(之后)
    NET本质论_读书笔记(1)
    WinDbg配置和使用基础(转)
    ASP.NET 2.0中CSS失效的问题总结(转)
    【下载】.NET Framework 源代码
    IL代码底层运行机制(转)
    asp.net水晶报表的一些问题
    Javascript 刷新框架及页面的方法总集
  • 原文地址:https://www.cnblogs.com/wujianming-110117/p/14182398.html
Copyright © 2011-2022 走看看