[Bug]: emulator.run运行时间有问题 #921

zhou-pz · 2024-11-24T12:21:29Z

Issue Type

Performance

Modules Involved

MPC protocol, SPU runtime

Have you reproduced the bug with SPU HEAD?

Yes

Have you searched existing issues?

Yes

SPU Version

spu 0.9.3dev20241009

OS Platform and Distribution

Ubuntu 22.04

Python Version

3.9

Compiler Version

GCC 11.2.1

Current Behavior?

下面代码中floyd_opt函数里最后一句代码 “batch_2 = batch_2.at[indices].set(batch_2_upper_triangle)” 导致程序跑得很慢，去掉后程序跑得很快，但是这是简单的赋值操作，不会带来这么大开销，是什么原因呢？输出日志和代码如下。

有batch_2 = batch_2.at[indices].set(batch_2_upper_triangle)时输出日志：

[2024-11-24 12:01:18.766] [info] [api.cc:172] [Profiling] SPU execution floyd_opt completed, input processing took 1.28e-06s, execution took 76.40721284s, output processing took 2.3442e-05s, total time 76.407237562s.
[2024-11-24 12:01:19.898] [info] [api.cc:220] HLO profiling: total time 76.390170644
[2024-11-24 12:01:19.898] [info] [api.cc:223] - pphlo.while, executed 100 times, duration 76.100738252s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - pphlo.custom_call: spu.gather, executed 100 times, duration 0.26589784s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - pphlo.concatenate, executed 196 times, duration 0.020669029s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - pphlo.slice, executed 394 times, duration 0.001173913s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - pphlo.reshape, executed 100 times, duration 0.000812756s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - pphlo.constant, executed 8 times, duration 0.000444215s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - pphlo.free, executed 792 times, duration 0.000280899s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - pphlo.convert, executed 1 times, duration 0.00015374s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:220] HAL profiling: total time 21.108059529
[2024-11-24 12:01:19.898] [info] [api.cc:223] - i_less, executed 1455400 times, duration 8.139717218s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - logical_not, executed 970200 times, duration 4.869683125s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - _mux, executed 485100 times, duration 4.673866196s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - _and, executed 1455300 times, duration 2.597090457s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - i_add, executed 485100 times, duration 0.827552852s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - seal, executed 1 times, duration 0.000149681s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:220] MPC profiling: total time 33.847582661000004
[2024-11-24 12:01:19.898] [info] [api.cc:223] - update_slice, executed 485100 times, duration 16.305582855s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - add_pp, executed 2910700 times, duration 2.964376123s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - extract_slice, executed 3396094 times, duration 2.585630088s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - reshape, executed 4366000 times, duration 2.440487449s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - negate_p, executed 2425600 times, duration 2.225913358s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - and_pp, executed 1455300 times, duration 1.603299254s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - broadcast, executed 2425603 times, duration 1.238344766s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - concatenate, executed 485296 times, duration 1.139825704s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - msb_p, executed 1455400 times, duration 1.137026837s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - add_aa, executed 970200 times, duration 0.70372609s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - make_p, executed 970200 times, duration 0.602565889s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - mul_ap, executed 485100 times, duration 0.353910826s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - transpose, executed 485100 times, duration 0.296557429s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - negate_a, executed 485100 times, duration 0.250191463s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:223] - p2a, executed 1 times, duration 0.00014453s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19.898] [info] [api.cc:233] Link details: total send bytes 0, recv bytes 0, send actions 0, recv actions 0
[2024-11-24 12:01:19,900] [ForkServerProcess-1] RunR: builtin_fetch_meta at node:0
[2024-11-24 12:01:19,902] [ForkServerProcess-1] RunR: builtin_fetch_object at node:0
[2024-11-24 12:01:19,903] [ForkServerProcess-2] RunR: builtin_fetch_object at node:1
[2024-11-24 12:01:19,905] [ForkServerProcess-3] RunR: builtin_fetch_object at node:2
[2024-11-24 12:01:19,907] Shutdown multiprocess cluster...

去掉batch_2 = batch_2.at[indices].set(batch_2_upper_triangle)时输出日志：

[2024-11-24 11:24:13.057] [info] [api.cc:172] [Profiling] SPU execution floyd_opt completed, input processing took 4.11e-07s, execution took 0.00039479s, output processing took 1.413e-06s, total time 0.000396614s.
[2024-11-24 11:24:13.057] [info] [api.cc:220] HLO profiling: total time 3.1869e-05
[2024-11-24 11:24:13.057] [info] [api.cc:223] - pphlo.constant, executed 1 times, duration 3.1869e-05s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 11:24:13.057] [info] [api.cc:220] HAL profiling: total time 0
[2024-11-24 11:24:13.057] [info] [api.cc:220] MPC profiling: total time 9.837e-06
[2024-11-24 11:24:13.057] [info] [api.cc:223] - broadcast, executed 1 times, duration 9.837e-06s, send bytes 0 recv bytes 0, send actions 0, recv actions 0
[2024-11-24 11:24:13.057] [info] [api.cc:233] Link details: total send bytes 0, recv bytes 0, send actions 0, recv actions 0
[2024-11-24 11:24:13,057] [ForkServerProcess-4] RunR: builtin_fetch_object at node:3
[2024-11-24 11:24:13,061] [ForkServerProcess-1] RunR: builtin_fetch_meta at node:0
[2024-11-24 11:24:13,062] [ForkServerProcess-1] RunR: builtin_fetch_object at node:0
[2024-11-24 11:24:13,065] [ForkServerProcess-2] RunR: builtin_fetch_object at node:1
[2024-11-24 11:24:13,066] [ForkServerProcess-3] RunR: builtin_fetch_object at node:2
[2024-11-24 11:24:13,067] Shutdown multiprocess cluster...

代码：

def floyd_opt(
    dist
):

    n = len(dist)
    batch_2 = dist

    for k in range(n):
        batch_2 = jnp.delete(batch_2, k, axis=0)
        col_k_without_dkk = batch_2[:, k]
        batch_2 = jnp.delete(batch_2, k, axis=1)
        dist_ik = jnp.zeros_like(batch_2)  
        dist_kj = jnp.zeros_like(batch_2)

        for i in range(n-1):
            if(i < k):
                dist_ik = dist_ik.at[i].set(jnp.full(n-1, dist[i][k]))
            else:
                dist_ik = dist_ik.at[i].set(jnp.full(n-1, dist[i+1][k]))

        for j in range(n-1):
            if(j < k):
                dist_kj = dist_kj.at[:, j].set(jnp.full(n-1, dist[k][j]))
            else:
                dist_kj = dist_kj.at[:, j].set(jnp.full(n-1, dist[k][j+1]))

        # 把上三角拿出来算
        indices = numpy.triu_indices(batch_2.shape[0], k=1)
        batch_2_upper_triangle = batch_2[indices]
        dist_ik_upper_triangle = dist_ik[indices]
        dist_kj_upper_triangle = dist_kj[indices]

        # # 把上三角放回去
        batch_2 = jnp.zeros_like(dist)
        batch_2 = batch_2.at[indices].set(batch_2_upper_triangle) ############ 这行代码对运行时间影响很大！！！！！
            
    return batch_2



def emul_cpz(mode: emulation.Mode.MULTIPROCESS):

    try:
        # bandwidth and latency only work for docker mode
        emulator = emulation.Emulator(
            emulation.CLUSTER_ABY3_3PC, mode, bandwidth=300, latency=20
        )
        emulator.up()

        # 设置样本数量和维度
        num_samples = 100
        Knn = np.random.rand(num_samples, num_samples)
        Knn = (Knn + Knn.T) / 2
        Knn[Knn == 0] = np.inf
        np.fill_diagonal(Knn, 0)
        Knn=emulator.seal(Knn)

        
        # floyd_opt
        shortest_paths_floyd= emulator.run(floyd_opt)(Knn)

    finally:
        emulator.down()

if __name__ == "__main__":
    emul_cpz(emulation.Mode.MULTIPROCESS)

Standalone code to reproduce the issue

如上

Relevant log output

如上

The text was updated successfully, but these errors were encountered:

deadlywing · 2024-11-25T02:41:33Z

hello，主要是因为如果那一行被comment，则编译器发现这个函数的return为常量0的array；
如果没被comment，才会按你定义的函数执行，至于为什么这么慢，主要是因为实现里的循环是O(n^2)的

zhou-pz · 2024-11-25T15:42:44Z

hello，我把里面的两个循环注释掉了时间没减少，还是需要78秒（dist和batch_2只是100*100的）

    for i in range(n-1):
        if(i < k):
            dist_ik = dist_ik.at[i].set(jnp.full(n-1, dist[i][k]))
        else:
            dist_ik = dist_ik.at[i].set(jnp.full(n-1, dist[i+1][k]))

    for j in range(n-1):
        if(j < k):
            dist_kj = dist_kj.at[:, j].set(jnp.full(n-1, dist[k][j]))
        else:
            dist_kj = dist_kj.at[:, j].set(jnp.full(n-1, dist[k][j+1]))

deadlywing · 2024-11-26T11:58:54Z

猜测主要原因还是 .at[].set[] 这个语句会产生大量的ops；虽然comm. free，但是cpu time还是比较多的。

建议想想有没有等价的方式可以实现。。

zhou-pz · 2024-11-28T03:53:36Z

.at[].set[]确实很慢，而且在emulator里格外慢：

(1) 我用emulator跑，需要75s。(2) 不用emulator跑（纯python程序），只需要10s。(3) 不用emulator跑，且把jax.numpy的.at[].set[]操作全部替换为numpy的标准索引赋值，只需要0.03s。

zhou-pz · 2024-11-28T03:57:16Z

请问如果赋值操作无法避免的话，spu支持其它更高效的赋值操作吗？

deadlywing · 2024-11-28T06:25:02Z

只从你提供的代码来看，看上去是要从dist这个矩阵中每次提取一些数据出来组成新的矩阵，不知道有没有可能通过一些特殊的index/gather方法来得到：可以看看jax.numpy.choose, jax.lax.gather等方法

github-actions · 2024-12-28T09:43:13Z

Stale issue message. Please comment to remove stale tag. Otherwise this issue will be closed soon.

github-actions bot added the no-issue-activity label Dec 28, 2024

github-actions bot closed this as not planned Won't fix, can't repro, duplicate, stale Jan 4, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Bug]: emulator.run运行时间有问题 #921

[Bug]: emulator.run运行时间有问题 #921

zhou-pz commented Nov 24, 2024

deadlywing commented Nov 25, 2024

zhou-pz commented Nov 25, 2024

deadlywing commented Nov 26, 2024

zhou-pz commented Nov 28, 2024

zhou-pz commented Nov 28, 2024

deadlywing commented Nov 28, 2024

github-actions bot commented Dec 28, 2024

[Bug]: emulator.run运行时间有问题 #921

[Bug]: emulator.run运行时间有问题 #921

Comments

zhou-pz commented Nov 24, 2024

Issue Type

Modules Involved

Have you reproduced the bug with SPU HEAD?

Have you searched existing issues?

SPU Version

OS Platform and Distribution

Python Version

Compiler Version

Current Behavior?

Standalone code to reproduce the issue

Relevant log output

deadlywing commented Nov 25, 2024

zhou-pz commented Nov 25, 2024

deadlywing commented Nov 26, 2024

zhou-pz commented Nov 28, 2024

zhou-pz commented Nov 28, 2024

deadlywing commented Nov 28, 2024

github-actions bot commented Dec 28, 2024