Inductor CPU 后端调试与性能分析

importtorch

deffoo1(x1, x2):
    a = torch.neg(x1)
    b = torch.maximum(x2, a)
    y = torch.cat([b], dim=0)
    return y

x1 = torch.randint(256, (1, 8), dtype=torch.uint8)
x2 = torch.randint(256, (8390, 8), dtype=torch.uint8)

compiled_foo1 = torch.compile(foo1)
result = compiled_foo1(x1, x2)

defneg1(x):
    return f"decltype({x})(-{x})"

exportKMP_BLOCKTIME=1
exportKMP_SETTINGS=1
exportKMP_AFFINITY=granularity=fine,compact,1,0
exportLD_PRELOAD=${CONDA_PREFIX:-"$(dirname$(whichconda))/../"}/lib/libiomp5.so:${CONDA_PREFIX:-"$(dirname$(whichconda))/../"}/lib/libjemalloc.so
exportMALLOC_CONF="oversize_threshold:1,background_thread:true,metadata_thp:auto,dirty_decay_ms:-1,muzzy_decay_ms:-1"
numactl-C0-31-m0pythonbench.py

# bench.py
fromtransformersimport MobileBertForQuestionAnswering
# Initialize an eager model
model = MobileBertForQuestionAnswering.from_pretrained("csarron/mobilebert-uncased-squad-v2")
seq_length = 128
bs = 128
vocab_size = model.config.vocab_size
input = torch.randint(0, vocab_size, (bs, seq_length), dtype=torch.int64)
input_dict = {"input_ids": input}

# Initialize the inductor model
compiled_model = torch.compile(model)
with torch.no_grad():
    compiled_model(**input_dict)

NUM_ITERS=50
importtimeit
with torch.no_grad():
    # warmup
    for _ in range(10):
        model(**input_dict)
    eager_t = timeit.timeit("model(**input_dict)", number=NUM_ITERS, globals=globals())

with torch.no_grad():
    # warmup
    for _ in range(10):
        compiled_model(**input_dict)
    inductor_t = timeit.timeit("compiled_model(**input_dict)", number=NUM_ITERS, globals=globals())
# print(f"eager use: {eager_t * 1000 / NUM_ITERS} ms/iter")
# print(f"inductor use: {inductor_t * 1000 / NUM_ITERS} ms/iter")
# print(f"speed up ratio: {eager_t / inductor_t}")

eageruse:802.1023553796113ms/iter
inductoruse:339.95180135127157ms/iter
speedupratio:2.359459053287382

fromtorch._inductorimport config
config.cpp.enable_kernel_profile = True

# bench.py
fromtorch.profilerimport profile, schedule, ProfilerActivity
RESULT_DIR = "./prof_trace"
my_schedule = schedule(
    skip_first=10,
    wait=5,
    warmup=5,
    active=1,
    repeat=5)

deftrace_handler(p):
    output = p.key_averages().table(sort_by="self_cpu_time_total", row_limit=20)
    # print(output)
    p.export_chrome_trace(f"{RESULT_DIR}/{p.step_num}.json")

for _ in range(10):
    model(**input_dict)  # compiled_model(**input_dict) to get inductor model profiling

total = 0
with profile(
    activities=[ProfilerActivity.CPU],
    schedule=my_schedule,
    on_trace_ready=trace_handler
) as p:
    for _ in range(50):
        model(**input_dict)  # compiled_model(**input_dict) to get inductor model profiling
        p.step()

*------------------------------------------------------------
NameCPUtotal%CPUtotal# of Calls
*------------------------------------------------------------
aten::addmm45.73%370.814ms362
aten::add19.89%161.276ms363
aten::copy_14.97%121.416ms488
aten::mul9.02%73.154ms194
aten::clamp_min8.81%71.444ms96
aten::bmm5.46%44.258ms48
ProfilerStep*100.00%810.920ms1
aten::div2.89%23.447ms24
aten::_softmax1.00%8.087ms24
aten::linear46.48%376.888ms362
aten::clone2.77%22.430ms98
aten::t0.31%2.502ms362
aten::view0.14%1.161ms850
aten::transpose0.17%1.377ms386
aten::index_select0.12%952.000us3
aten::expand0.12%986.000us458
aten::matmul8.31%67.420ms48
aten::cat0.09%703.000us1
aten::as_strided0.08%656.000us963
aten::relu8.86%71.864ms96
*------------------------------------------------------------
SelfCPUtimetotal:810.920ms

*----------------------------------------------------------------------------------
NameCPUtotal%CPUtotal# of Calls
*----------------------------------------------------------------------------------
mkl::_mkl_linear68.79%231.573ms362
aten::bmm8.02%26.992ms48
ProfilerStep*100.00%336.642ms1
graph_0_cpp_fused_constant_pad_nd_embedding_00.27%915.000us1
aten::empty0.27%911.000us362
graph_0_cpp_fused__mkl_linear_add_mul_relu_1510.27%901.000us1
graph_0_cpp_fused__mkl_linear_add_mul_relu_2260.27%899.000us1
graph_0_cpp_fused__mkl_linear_add_mul_relu_3610.27%898.000us1
graph_0_cpp_fused__mkl_linear_add_mul_relu_1210.27%895.000us1
graph_0_cpp_fused__mkl_linear_add_mul_relu_310.27%893.000us1
graph_0_cpp_fused__mkl_linear_add_mul_relu_760.26%892.000us1
graph_0_cpp_fused__mkl_linear_add_mul_relu_2560.26%892.000us1
graph_0_cpp_fused__mkl_linear_add_mul_relu_3460.26%892.000us1
graph_0_cpp_fused__mkl_linear_add_mul_relu_2410.26%891.000us1
graph_0_cpp_fused__mkl_linear_add_mul_relu_3160.26%891.000us1
graph_0_cpp_fused__mkl_linear_add_mul_relu_910.26%890.000us1
graph_0_cpp_fused__mkl_linear_add_mul_relu_1060.26%890.000us1
graph_0_cpp_fused__mkl_linear_add_mul_relu_2110.26%890.000us1
graph_0_cpp_fused__mkl_linear_add_mul_relu_610.26%889.000us1
graph_0_cpp_fused__mkl_linear_add_mul_relu_2860.26%889.000us1
*----------------------------------------------------------------------------------
SelfCPUtimetotal:336.642ms

cpp_fused__mkl_linear_add_mul_relu_151 = async_compile.cpp('''
#include <ATen/record_function.h>
#include "/tmp/torchinductor_root/lr/clrlgu27q4ggd472umdzwsu6qcpqxcuusjxqvx2hwitjbujiiz7z.h"
extern "C" void kernel(float* in_out_ptr0,
                       const float* in_ptr0,
                       const float* in_ptr1,
                       const float* in_ptr2,
                       const float* in_ptr3)
{
    RECORD_FUNCTION("graph_0_cpp_fused__mkl_linear_add_mul_relu_151", c10::ArrayRef<c10::IValue>({}));
    #pragma omp parallel num_threads(32)
    {
        {
            #pragma omp for
            for(long i0=static_cast<long>(0L); i0<static_cast<long>(16384L); i0+=static_cast<long>(1L))
            {
                for(long i1=static_cast<long>(0L); i1<static_cast<long>(512L); i1+=static_cast<long>(8L))
                {
                    auto tmp0 = at::vec::Vectorized<float>::loadu(in_ptr0 + static_cast<long>(i1 + (512L*i0)));
                    auto tmp1 = at::vec::Vectorized<float>::loadu(in_ptr1 + static_cast<long>(i1));
                    auto tmp3 = at::vec::Vectorized<float>::loadu(in_out_ptr0 + static_cast<long>(i1 + (512L*i0)));
                    auto tmp5 = at::vec::Vectorized<float>::loadu(in_ptr2 + static_cast<long>(i1));
                    auto tmp7 = at::vec::Vectorized<float>::loadu(in_ptr3 + static_cast<long>(i1));
                    auto tmp2 = tmp0 + tmp1;
                    auto tmp4 = tmp2 + tmp3;
                    auto tmp6 = tmp4 * tmp5;
                    auto tmp8 = tmp6 + tmp7;
                    tmp8.store(in_out_ptr0 + static_cast<long>(i1 + (512L*i0)));
                }
            }
        }
    }
}''')

# bench.py
deffunc(arg_0, arg_1, arg_2, arg_3, arg_4):
    add_0 = arg_0 + arg_1
    add_1 = add_0 + arg_2
    mul_1 = add_1 * arg_3
    add_2 = mul_1 + arg_4
    arg_2 = add_2
    return arg_2

arg_0 = torch.rand(16384, 512)
arg_1 = torch.rand(1, 512)
arg_2 = torch.zeros(16384, 512)
arg_3 = torch.rand(1, 512)
arg_4 = torch.rand(1, 512)

input = (arg_0, arg_1, arg_2, arg_3, arg_4)
inductor_func = torch.compile(func)
with torch.no_grad():
    inductor_func(*input)

importtimeit
NUM_ITERS=100
with torch.no_grad():
    # warmup
    for _ in range(10):
        func(*input)
    eager_t = timeit.timeit("func(*input)", number=NUM_ITERS, globals=globals())

with torch.no_grad():
    # warmup
    for _ in range(10):
        inductor_func(*input)
    inductor_t = timeit.timeit("inductor_func(*input)", number=NUM_ITERS, globals=globals())
# print(f"eager use: {eager_t * 1000 / NUM_ITERS} ms/iter")
# print(f"inductor use: {inductor_t * 1000 / NUM_ITERS} ms/iter")
# print(f"speed up ratio: {eager_t / inductor_t}")

eageruse:5.780875144992024ms/iter
inductoruse:0.9588955780491233ms/iter
speedupratio:6.0286805751604735

下载 Python 源代码: inductor_debug_cpu.py

下载 Jupyter 笔记本: inductor_debug_cpu.ipynb