optimize ggml_ext_chunk

[leejet]

No description provided.

[leejet]

Perhaps we can remove ggml_cont here, but I haven’t fully verified it yet, so I’ll keep ggml_cont here for now.

[stduhpf]

Wouldn't something like this work in all cases? It seems to work for unet models at least. The code is simpler and it seems (slightly) faster than using permutations.

__STATIC_INLINE__ std::vector<struct ggml_tensor*> ggml_ext_chunk(struct ggml_context* ctx,
                                                                  struct ggml_tensor* x,
                                                                  int num,
                                                                  int64_t dim,
                                                                  bool cont = true) {
    GGML_ASSERT(dim >= 0 && dim < 4);
    GGML_ASSERT(x->ne[dim] % num == 0);

    std::vector<struct ggml_tensor*> chunks;
    int64_t chunk_size  = x->ne[dim] / num;
    int64_t stride      = chunk_size * x->nb[dim];
    int64_t chunk_ne[4] = {x->ne[0], x->ne[1], x->ne[2], x->ne[3]};
    chunk_ne[dim]       = chunk_size;
    for (int i = 0; i < num; i++) {
        auto chunk = ggml_view_4d(
            ctx, x,
            chunk_ne[0], chunk_ne[1], chunk_ne[2], chunk_ne[3],
            x->nb[1], x->nb[2], x->nb[3], stride * i);
        if (cont) {
            chunk = ggml_cont(ctx, chunk);
        }
        chunks.push_back(chunk);
    }

    return chunks;
}

[stduhpf]

Maybe it's not really faster, seems to be withing margin of error for run-to-run variations.

[stduhpf]

Also it doesn't cause #1080 (comment)

[leejet]

@stduhpf This change looks like a simpler optimization. Could you open a separate PR for it? I’ll close this PR.

[wbruna]

0835e5c broke sd1.5:

master-408	0835e5c

[stduhpf]

@wbruna, Oh, you're right, I was only looking at the speed.

[daniandtheweb]

0835e5c broke sd1.5:

Same on SDXL.

[wbruna]

Testing each version on SD1.5: when compared with 59ebdf0, #1079 seems almost as fast on Vulkan, and around 9% slower on ROCm. The ggml_ext_chunk suggested above is ~3-4% slower on both:

version	vulkan	rocm
59ebdf0	2.65s/it	2.34s/it
347710f (and current master)	3.65s/it	3.44s/it
ggml_ext_chunk above	2.75s/it	2.41s/it
#1079	2.69s/it	2.54s/it

[leejet]

0835e5c broke sd1.5:

master-408 0835e5c

It looks like the implementations of the CUDA backend and the Vulkan backend are a bit different. I was able to reproduce it with the Vulkan backend as well, but everything works fine with the CUDA backend.