Support cuTensor?

[wxj6000]

I haven't see the complete documentation of this project. But the binding for cuTensor is not found in the repo. I wonder if there is any plan for supporting general tensor contractions?

[leofang]

Hi @wxj6000 thanks for reaching out and opening the very first issue in this repo! 😄

cuTENSOR support is on our roadmap but there is no ETA yet. I am sure you are aware that both CuPy and cuQuantum Python support cuTENSOR. We would be very curious to learn what use cases that you have in mind are not already covered by these projects. It would help us prioritize for new features.

FYI, the documentation will be live next week.

[wxj6000]

@leofang Thanks. I was trying to understand the scope of the project and show my interest.

CuPy and cuQuantum Python do have the support for cuTensor. My collaborator and I am considering the kernel fusion of two tensor contraction kernels. For example,

c = contract('ijk,ko->ijo', a, b)
d = contract('ijo,ilo->jl', c, c)

[yangcal]

Hello @wxj6000,

With nvmath beta 7, we now have a [nvmath.tensor module] for linux platform(https://docs.nvidia.com/cuda/nvmath-python/latest/host-apis/tensor/index.html) that support both binary and ternary contraction. Bindings are provided as well under nvmath.bindings and cuTensor>2.3.1 is required.

Feel free to experiment with our tensor module or refer to these simple examples. We would love to hear any feedbacks.

[wxj6000]

Hi @yangcal , thank you for following this up. The new feature looks interesting. Recently, I had too many tasks to complete. I will test the new functions this week and post the benchmarks here.

[wxj6000]

Hi the nvmath-python team, I recently tried nvmath.tensor.TernaryContraction and collected benchmarks (see the attachment). And I hope they are helpful.

I tried two typical contraction patterns, 'ig,ij,jg->g' and 'g,ig,jg->ij' which are used in quantum chemistry. I fixed g dimension as 16384. i and j dimensions vary. nvmath is consistently faster than cp.einsum for the first one, but slower in the later one. I expected the contraction is memory-bound in both cases when i and j is small. nvmath's fused contraction can benefit, which is true for the first case. Is there any fundamental difficulty for the second contraction pattern?

The overhead of creating a plan is significant, ~200 ms, which is much more expensive the computation. The stateful API is really helpful to reduce the cost of overhead.

Here are the benchmarks on A100-80G.

[samaid]

Hi @wxj6000,

Appreciate your experiments. We are taking care of the overheads (not only in tensor contractions but overall across the entire library). Hope it will become better in the coming release.

Let us also take a look at your test case. At first glance, I cannot reasonably explain why g,ig, jg->ij is slower than CuPy.

[wxj6000]

Hi @wxj6000,

Appreciate your experiments. We are taking care of the overheads (not only in tensor contractions but overall across the entire library). Hope it will become better in the coming release.

This would be very helpful!

Let us also take a look at your test case. At first glance, I cannot reasonably explain why g,ig, jg->ij is slower than CuPy.

It might be on the kernel selection? The dimension to be reduced,g dimension, is relatively large. The issue could also be on me. So I posted the script so that your team can reproduce.

import cupy as cp
import nvmath
import time

i_dim = 512
j_dim = 512
g_dim = 16384

a = cp.random.rand(i_dim, g_dim)
b = cp.random.rand(j_dim, g_dim)
c = cp.random.rand(i_dim, j_dim)
g = cp.random.rand(g_dim)

# FLOPs calculation
flops_contraction1 = g_dim * (3 * i_dim * j_dim - 1)
flops_contraction2 = flops_contraction1
flops_contraction3 = i_dim * j_dim * (3 * g_dim - 1)

# Time nvmath.tensor.ternary_contraction
# ======================================
# Warmup step to compile kernels, etc.
num_runs = 10

# Timed run
nvmath_times = []
with nvmath.tensor.TernaryContraction("ig,ij,jg->g", a, c, b) as contraction:
    contraction.plan()
    a[:] = cp.random.rand(i_dim, g_dim)
    b[:] = cp.random.rand(j_dim, g_dim)
    c[:] = cp.random.rand(i_dim, j_dim)
    start_event = cp.cuda.Event()
    end_event = cp.cuda.Event()
    start_event.record()
    r1 = contraction.execute()
    end_event.record()
    end_event.synchronize()
    warmup_time_nvmath = cp.cuda.get_elapsed_time(start_event, end_event)
    print(f"nvmath.tensor.ternary_contraction (ig,ij,jg->g, warmup) took: {warmup_time_nvmath:.4f} ms")
    for _ in range(num_runs):
        a[:] = cp.random.rand(i_dim, g_dim)
        b[:] = cp.random.rand(j_dim, g_dim)
        c[:] = cp.random.rand(i_dim, j_dim)
        start_event = cp.cuda.Event()
        end_event = cp.cuda.Event()
        start_event.record()
        result = contraction.execute()
        end_event.record()
        end_event.synchronize()
        nvmath_times.append(cp.cuda.get_elapsed_time(start_event, end_event))
t_nvmath = sum(nvmath_times) / num_runs
gflops_nvmath = (flops_contraction1 / (t_nvmath / 1000)) / 1e9
print(f"nvmath.tensor.ternary_contraction (ig,ij,jg->g, timed, average over {num_runs} runs) took: {t_nvmath:.4f} ms, {gflops_nvmath:.4f} GFLOPS")


# Time cp.einsum split into two contractions
# ==========================================
# Warmup step
start_event = cp.cuda.Event()
end_event = cp.cuda.Event()
start_event.record()
intermediate = cp.einsum("ig,ij->jg", a, c)
rho_split = cp.einsum("jg,jg->g", intermediate, b)
end_event.record()
end_event.synchronize()
warmup_time_einsum_split = cp.cuda.get_elapsed_time(start_event, end_event)
print(f"cp.einsum (ig,ij->jg and jg,jg->g, split, warmup) took: {warmup_time_einsum_split:.4f} ms")

# Timed run
einsum_split_times = []
for _ in range(num_runs):
    start_event = cp.cuda.Event()
    end_event = cp.cuda.Event()
    start_event.record()
    intermediate = cp.einsum("ig,ij->jg", a, c)
    rho_split = cp.einsum("jg,jg->g", intermediate, b)
    end_event.record()
    end_event.synchronize()
    einsum_split_times.append(cp.cuda.get_elapsed_time(start_event, end_event))
t_einsum_split = sum(einsum_split_times) / num_runs
gflops_einsum_split = (flops_contraction2 / (t_einsum_split / 1000)) / 1e9
print(f"cp.einsum (ig,ij->jg and jg,jg->g, split, timed, average over {num_runs} runs) took: {t_einsum_split:.4f} ms, {gflops_einsum_split:.4f} GFLOPS")


assert cp.allclose(result, rho_split)

# New contraction pattern: g,ig,jg->ij
# ====================================
# Time nvmath.tensor.ternary_contraction
# Warmup step
start_event = cp.cuda.Event()
end_event = cp.cuda.Event()
start_event.record()
result_2 = nvmath.tensor.ternary_contraction("g,ig,jg->ij", g, a, b)
end_event.record()
end_event.synchronize()
warmup_time_nvmath_2 = cp.cuda.get_elapsed_time(start_event, end_event)
print(f"nvmath.tensor.ternary_contraction (g,ig,jg->ij, warmup) took: {warmup_time_nvmath_2:.4f} ms")

# Timed run
nvmath_times_2 = []
with nvmath.tensor.TernaryContraction("g,ig,jg->ij", g, a, b) as contraction:
    contraction.plan()
    g[:] = cp.random.rand(g_dim)
    a[:] = cp.random.rand(i_dim, g_dim)
    b[:] = cp.random.rand(j_dim, g_dim)
    c[:] = cp.random.rand(i_dim, j_dim)
    start_event = cp.cuda.Event()
    end_event = cp.cuda.Event()
    start_event.record()
    result_2 = contraction.execute()
    end_event.record()
    end_event.synchronize()
    warmup_time_nvmath_2 = cp.cuda.get_elapsed_time(start_event, end_event)
    print(f"nvmath.tensor.ternary_contraction (g,ig,jg->ij, warmup) took: {warmup_time_nvmath_2:.4f} ms")
    for _ in range(num_runs):
        g[:] = cp.random.rand(g_dim)
        a[:] = cp.random.rand(i_dim, g_dim)
        b[:] = cp.random.rand(j_dim, g_dim)
        c[:] = cp.random.rand(i_dim, j_dim)
        start_event = cp.cuda.Event()
        end_event = cp.cuda.Event()
        start_event.record()
        result_2 = contraction.execute()
        end_event.record()
        end_event.synchronize()
        nvmath_times_2.append(cp.cuda.get_elapsed_time(start_event, end_event))
t_nvmath_2 = sum(nvmath_times_2) / num_runs
gflops_nvmath_2 = (flops_contraction3 / (t_nvmath_2 / 1000)) / 1e9
print(f"nvmath.tensor.ternary_contraction (g,ig,jg->ij, timed, average over {num_runs} runs) took: {t_nvmath_2:.4f} ms, {gflops_nvmath_2:.4f} GFLOPS")


# Time cp.einsum for 'g,ig->ig' and 'ig,jg->ij' together
# =======================================================
# Warmup step
start_event = cp.cuda.Event()
end_event = cp.cuda.Event()
start_event.record()
intermediate = cp.einsum("g,ig->ig", g, a)
result_einsum_2_combined = cp.einsum("ig,jg->ij", intermediate, b)
end_event.record()
end_event.synchronize()
warmup_time_einsum_combined = cp.cuda.get_elapsed_time(start_event, end_event)
print(f"cp.einsum (g,ig->ig and ig,jg->ij, combined, warmup) took: {warmup_time_einsum_combined:.4f} ms")

# Timed run
einsum_combined_times = []
for _ in range(num_runs):
    start_event = cp.cuda.Event()
    end_event = cp.cuda.Event()
    start_event.record()
    intermediate = cp.einsum("g,ig->ig", g, a)
    result_einsum_2_combined = cp.einsum("ig,jg->ij", intermediate, b)
    end_event.record()
    end_event.synchronize()
    einsum_combined_times.append(cp.cuda.get_elapsed_time(start_event, end_event))
t_einsum_combined = sum(einsum_combined_times) / num_runs
# gflops calculation for g,ig,jg->ij: i_dim * j_dim * (3 * g_dim - 1)
gflops_einsum_combined = (flops_contraction3 / (t_einsum_combined / 1000)) / 1e9
print(f"cp.einsum (g,ig->ig and ig,jg->ij, combined, timed, average over {num_runs} runs) took: {t_einsum_combined:.4f} ms, {gflops_einsum_combined:.4f} GFLOPS")


# check correctness
einsum_result_2 = cp.einsum("g,ig->ig", g, a)
einsum_result_2 = cp.einsum("ig,jg->ij", einsum_result_2, b)
assert cp.allclose(result_2, einsum_result_2)
assert cp.allclose(result_2, result_einsum_2_combined)

[yangcal]

Hello @wxj6000

We were able to reproduce the issue, which appears to trace back to one of cutensor's kernel. We have filed an internal ticket to resolve the perf issue. Thanks for sharing the experience.