Is is possible to parallelize sparse-dense matrix mul on gpus and tpus?

[dontknow]

Hi.

a) does the following run in parallel on all available gpus??

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    strategy.run(tf.sparse.sparse_dense_matmul, args=(sparseMat, denseVec)), axis=None)

b) if not, is it possible to run sparse, dense matrix multiplication in parallel in tensorflow on all gpus???

c) does the following run in parallel on all available tpus??

strategy = tf.distribute.TPUStrategy(resolver)
with strategy.scope():
    strategy.run(tf.sparse.sparse_dense_matmul, args=(sparseMat, denseVec)), axis=None)

d) if not, is it possible to run sparse dense matrix multiplication in parallel in tensorflow on all tpus???

Thank you.
Don’t Know.

[Aniket_Dubey]

Hi @dontknow ,
Welcome to the Google AI forum .

@dontknow

does the following run in parallel on all available gpus

The provided code using tf.distribute.MirroredStrategy() does not automatically parallelize tf.sparse.sparse_dense_matmul across GPUs. The strategy.run method does distribute the computation to each GPU, but it does not handle the distribution of a single sparse-dense matrix multiplication operation across multiple GPUs. For parallel execution across GPUs, you would need to manually distribute data and handle aggregation.

@dontknow

if not, is it possible to run sparse, dense matrix multiplication in parallel in tensorflow on all gpus???

To perform sparse-dense matrix multiplication in parallel on all GPUs, you need to manually implement a parallelization strategy. This typically involves distributing the input data across GPUs, performing computations in parallel, and then aggregating the results. TensorFlow’s standard API does not directly support this for sparse-dense operations across multiple GPUs.

@dontknow

does the following run in parallel on all available tpus

Using tf.distribute.TPUStrategy() with strategy.run does not automatically parallelize tf.sparse.sparse_dense_matmul across all TPU cores. While strategy.run distributes computation across TPU cores, the tf.sparse.sparse_dense_matmul function itself does not natively support multi-core parallel execution in this context without further optimization.

@dontknow

d) if not, is it possible to run sparse dense matrix multiplication in parallel in tensorflow on all tpus???

Sparse-dense matrix multiplication can be parallelized on TPUs, but it generally requires custom implementation or optimization to leverage TPU cores efficiently. TensorFlow provides support for sparse operations on TPUs, but distributing a sparse-dense matrix multiplication across TPU cores may need additional steps or use of TPU-specific operations and libraries.

However, there’s an important caveat for both GPU and TPU implementations: The tf.sparse.sparse_dense_matmul operation itself may not be optimized for parallel execution on multiple devices.

Thank You .

[Sarah_More]

**strategy = emphasized texttf.distribute.TPUStrategy(resolver)
with strategy.scope():
strategy.run(tf.sparse.sparse_dense_matmul, args=(sparseMat, denseVec)), ax

is=None)**