<html><head><meta http-equiv="content-type" content="text/html; charset=utf-8"></head><body style="overflow-wrap: break-word; -webkit-nbsp-mode: space; line-break: after-white-space;"><div dir="auto" style="overflow-wrap: break-word; -webkit-nbsp-mode: space; line-break: after-white-space;">Thanks a lot Pierre!<div><br></div><div>I managed to solve my problem for now using the (less scalable) solution that you provided. I also opened up an issue about the missing MatMPIAIJGetLocalMat() petsc4py binding here <a href="https://gitlab.com/petsc/petsc/-/issues/1443">https://gitlab.com/petsc/petsc/-/issues/1443</a> and if no action is taken, I might take care of it myself as well, soon.</div><div><br></div><div>For the sake of completeness, and to help any potential PETSc user that might run into the same issue, here is my final code that works nicely in parallel.</div><div><br></div><div><div style="color: rgb(29, 33, 38); background-color: rgb(255, 255, 255); font-family: Menlo, Monaco, "Courier New", monospace; line-height: 18px; white-space: pre;"><div><span style="color: rgb(10, 66, 131);">"""Experimenting with PETSc mat-mat multiplication"""</span></div><br><div><span style="color: rgb(168, 3, 20);">import</span> <span style="color: rgb(72, 31, 148);">numpy</span> <span style="color: rgb(168, 3, 20);">as</span> <span style="color: rgb(72, 31, 148);">np</span></div><div><span style="color: rgb(168, 3, 20);">from</span> <span style="color: rgb(72, 31, 148);">firedrake</span> <span style="color: rgb(168, 3, 20);">import</span> <span style="color: rgb(13, 92, 181);">COMM_WORLD</span></div><div><span style="color: rgb(168, 3, 20);">from</span> <span style="color: rgb(72, 31, 148);">firedrake</span>.<span style="color: rgb(72, 31, 148);">petsc</span> <span style="color: rgb(168, 3, 20);">import</span> <span style="color: rgb(72, 31, 148);">PETSc</span></div><div><span style="color: rgb(168, 3, 20);">from</span> <span style="color: rgb(72, 31, 148);">numpy</span>.<span style="color: rgb(72, 31, 148);">testing</span> <span style="color: rgb(168, 3, 20);">import</span> <span style="color: rgb(72, 31, 148);">assert_array_almost_equal</span></div><br><div><span style="color: rgb(164, 176, 188);"># import pdb</span></div><br><div>nproc <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">COMM_WORLD</span>.size</div><div>rank <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">COMM_WORLD</span>.rank</div><br><br><div><span style="color: rgb(168, 3, 20);">def</span> <span style="color: rgb(72, 31, 148);">Print</span>(<span style="color: rgb(166, 61, 0);">x</span>: <span style="color: rgb(72, 31, 148);">str</span>):</div><div> <span style="color: rgb(10, 66, 131);">"""Prints the string only on the root process</span></div><br><div><span style="color: rgb(10, 66, 131);"> Args:</span></div><div><span style="color: rgb(10, 66, 131);"> x (str): String to be printed</span></div><div><span style="color: rgb(10, 66, 131);"> """</span></div><div> <span style="color: rgb(72, 31, 148);">PETSc</span>.Sys.Print(<span style="color: rgb(166, 61, 0);">x</span>)</div><br><br><div><span style="color: rgb(168, 3, 20);">def</span> <span style="color: rgb(72, 31, 148);">print_mat_info</span>(<span style="color: rgb(166, 61, 0);">mat</span>, <span style="color: rgb(166, 61, 0);">name</span>):</div><div> <span style="color: rgb(10, 66, 131);">"""Prints the matrix information</span></div><br><div><span style="color: rgb(10, 66, 131);"> Args:</span></div><div><span style="color: rgb(10, 66, 131);"> mat (PETSc mat): PETSc matrix</span></div><div><span style="color: rgb(10, 66, 131);"> name (string): Name of the matrix</span></div><div><span style="color: rgb(10, 66, 131);"> """</span></div><div> <span style="color: rgb(72, 31, 148);">Print</span>(<span style="color: rgb(168, 3, 20);">f</span><span style="color: rgb(10, 66, 131);">"MATRIX </span><span style="color: rgb(13, 92, 181);">{</span><span style="color: rgb(166, 61, 0);">name</span><span style="color: rgb(13, 92, 181);">}</span><span style="color: rgb(10, 66, 131);"> [</span><span style="color: rgb(13, 92, 181);">{</span><span style="color: rgb(166, 61, 0);">mat</span>.getSize()[<span style="color: rgb(13, 92, 181);">0</span>]<span style="color: rgb(13, 92, 181);">}</span><span style="color: rgb(10, 66, 131);">x</span><span style="color: rgb(13, 92, 181);">{</span><span style="color: rgb(166, 61, 0);">mat</span>.getSize()[<span style="color: rgb(13, 92, 181);">1</span>]<span style="color: rgb(13, 92, 181);">}</span><span style="color: rgb(10, 66, 131);">]"</span>)</div><div> <span style="color: rgb(164, 176, 188);"># print(f"For rank {rank} local {name}: {mat.getSizes()}")</span></div><div> <span style="color: rgb(72, 31, 148);">Print</span>(<span style="color: rgb(166, 61, 0);">mat</span>.getType())</div><div> <span style="color: rgb(166, 61, 0);">mat</span>.view()</div><div> <span style="color: rgb(72, 31, 148);">Print</span>(<span style="color: rgb(10, 66, 131);">""</span>)</div><div> <span style="color: rgb(13, 92, 181);">COMM_WORLD</span>.Barrier()</div><div> <span style="color: rgb(72, 31, 148);">Print</span>(<span style="color: rgb(10, 66, 131);">""</span>)</div><br><br><div><span style="color: rgb(168, 3, 20);">def</span> <span style="color: rgb(72, 31, 148);">create_petsc_matrix_seq</span>(<span style="color: rgb(166, 61, 0);">input_array</span>):</div><div> <span style="color: rgb(10, 66, 131);">"""Building a sequential PETSc matrix from an array</span></div><br><div><span style="color: rgb(10, 66, 131);"> Args:</span></div><div><span style="color: rgb(10, 66, 131);"> input_array (np array): Input array</span></div><br><div><span style="color: rgb(10, 66, 131);"> Returns:</span></div><div><span style="color: rgb(10, 66, 131);"> seq mat: PETSc matrix</span></div><div><span style="color: rgb(10, 66, 131);"> """</span></div><div> <span style="color: rgb(168, 3, 20);">assert</span> <span style="color: rgb(72, 31, 148);">len</span>(<span style="color: rgb(166, 61, 0);">input_array</span>.shape) <span style="color: rgb(168, 3, 20);">==</span> <span style="color: rgb(13, 92, 181);">2</span></div><br><div> m, n <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">input_array</span>.shape</div><div> matrix <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.Mat().createAIJ(<span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>(m, n), <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(72, 31, 148);">PETSc</span>.<span style="color: rgb(13, 92, 181);">COMM_SELF</span>)</div><div> matrix.setUp()</div><br><div> matrix.setValues(<span style="color: rgb(72, 31, 148);">range</span>(m), <span style="color: rgb(72, 31, 148);">range</span>(n), <span style="color: rgb(166, 61, 0);">input_array</span>, <span style="color: rgb(166, 61, 0);">addv</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">False</span>)</div><br><div> <span style="color: rgb(164, 176, 188);"># Assembly the matrix to compute the final structure</span></div><div> matrix.assemblyBegin()</div><div> matrix.assemblyEnd()</div><br><div> <span style="color: rgb(168, 3, 20);">return</span> matrix</div><br><br><div><span style="color: rgb(168, 3, 20);">def</span> <span style="color: rgb(72, 31, 148);">create_petsc_matrix</span>(<span style="color: rgb(166, 61, 0);">input_array</span>, <span style="color: rgb(166, 61, 0);">sparse</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">True</span>):</div><div> <span style="color: rgb(10, 66, 131);">"""Create a PETSc matrix from an input_array</span></div><br><div><span style="color: rgb(10, 66, 131);"> Args:</span></div><div><span style="color: rgb(10, 66, 131);"> input_array (np array): Input array</span></div><div><span style="color: rgb(10, 66, 131);"> partition_like (PETSc mat, optional): Petsc matrix. Defaults to None.</span></div><div><span style="color: rgb(10, 66, 131);"> sparse (bool, optional): Toggle for sparese or dense. Defaults to True.</span></div><br><div><span style="color: rgb(10, 66, 131);"> Returns:</span></div><div><span style="color: rgb(10, 66, 131);"> PETSc mat: PETSc matrix</span></div><div><span style="color: rgb(10, 66, 131);"> """</span></div><div> <span style="color: rgb(164, 176, 188);"># Check if input_array is 1D and reshape if necessary</span></div><div> <span style="color: rgb(168, 3, 20);">assert</span> <span style="color: rgb(72, 31, 148);">len</span>(<span style="color: rgb(166, 61, 0);">input_array</span>.shape) <span style="color: rgb(168, 3, 20);">==</span> <span style="color: rgb(13, 92, 181);">2</span>, <span style="color: rgb(10, 66, 131);">"Input array should be 2-dimensional"</span></div><div> global_rows, global_cols <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">input_array</span>.shape</div><div><br></div><div> size <span style="color: rgb(168, 3, 20);">=</span> ((<span style="color: rgb(13, 92, 181);">None</span>, global_rows), (global_cols, global_cols))</div><br><div> <span style="color: rgb(164, 176, 188);"># Create a sparse or dense matrix based on the 'sparse' argument</span></div><div> <span style="color: rgb(168, 3, 20);">if</span> <span style="color: rgb(166, 61, 0);">sparse</span>:</div><div> matrix <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.Mat().createAIJ(<span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>size, <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">COMM_WORLD</span>)</div><div> <span style="color: rgb(168, 3, 20);">else</span>:</div><div> matrix <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.Mat().createDense(<span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>size, <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">COMM_WORLD</span>)</div><div> matrix.setUp()</div><br><div> local_rows_start, local_rows_end <span style="color: rgb(168, 3, 20);">=</span> matrix.getOwnershipRange()</div><br><div> <span style="color: rgb(168, 3, 20);">for</span> counter, i <span style="color: rgb(168, 3, 20);">in</span> <span style="color: rgb(72, 31, 148);">enumerate</span>(<span style="color: rgb(72, 31, 148);">range</span>(local_rows_start, local_rows_end)):</div><div> <span style="color: rgb(164, 176, 188);"># Calculate the correct row in the array for the current process</span></div><div> row_in_array <span style="color: rgb(168, 3, 20);">=</span> counter <span style="color: rgb(168, 3, 20);">+</span> local_rows_start</div><div> matrix.setValues(</div><div> i, <span style="color: rgb(72, 31, 148);">range</span>(global_cols), <span style="color: rgb(166, 61, 0);">input_array</span>[row_in_array, :], <span style="color: rgb(166, 61, 0);">addv</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">False</span></div><div> )</div><br><div> <span style="color: rgb(164, 176, 188);"># Assembly the matrix to compute the final structure</span></div><div> matrix.assemblyBegin()</div><div> matrix.assemblyEnd()</div><br><div> <span style="color: rgb(168, 3, 20);">return</span> matrix</div><br><br><div><span style="color: rgb(168, 3, 20);">def</span> <span style="color: rgb(72, 31, 148);">get_local_submatrix</span>(<span style="color: rgb(166, 61, 0);">A</span>):</div><div> <span style="color: rgb(10, 66, 131);">"""Get the local submatrix of A</span></div><br><div><span style="color: rgb(10, 66, 131);"> Args:</span></div><div><span style="color: rgb(10, 66, 131);"> A (mpi PETSc mat): partitioned PETSc matrix</span></div><br><div><span style="color: rgb(10, 66, 131);"> Returns:</span></div><div><span style="color: rgb(10, 66, 131);"> seq mat: PETSc matrix</span></div><div><span style="color: rgb(10, 66, 131);"> """</span></div><div> local_rows_start, local_rows_end <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">A</span>.getOwnershipRange()</div><div> local_rows <span style="color: rgb(168, 3, 20);">=</span> local_rows_end <span style="color: rgb(168, 3, 20);">-</span> local_rows_start</div><div> comm <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">A</span>.getComm()</div><div> rows <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.IS().createStride(</div><div> local_rows, <span style="color: rgb(166, 61, 0);">first</span><span style="color: rgb(168, 3, 20);">=</span>local_rows_start, <span style="color: rgb(166, 61, 0);">step</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">1</span>, <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span>comm</div><div> )</div><div> _, k <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">A</span>.getSize() <span style="color: rgb(164, 176, 188);"># Get the number of columns (k) from A's size</span></div><div> cols <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.IS().createStride(k, <span style="color: rgb(166, 61, 0);">first</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">0</span>, <span style="color: rgb(166, 61, 0);">step</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">1</span>, <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span>comm)</div><br><div> <span style="color: rgb(164, 176, 188);"># Getting the local submatrix</span></div><div> <span style="color: rgb(164, 176, 188);"># </span><span style="color: rgb(168, 3, 20);">TODO</span><span style="color: rgb(164, 176, 188);">: To be replaced by MatMPIAIJGetLocalMat() in the future (see petsc-users mailing list). There is a missing petsc4py binding, need to add it myself (and please create a merge request)</span></div><div> A_local <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">A</span>.createSubMatrices(rows, cols)[<span style="color: rgb(13, 92, 181);">0</span>]</div><div> <span style="color: rgb(168, 3, 20);">return</span> A_local</div><br><br><div><span style="color: rgb(168, 3, 20);">def</span> <span style="color: rgb(72, 31, 148);">multiply_sequential_matrices</span>(<span style="color: rgb(166, 61, 0);">A_seq</span>, <span style="color: rgb(166, 61, 0);">B_seq</span>):</div><div> <span style="color: rgb(10, 66, 131);">"""Multiply 2 sequential matrices</span></div><br><div><span style="color: rgb(10, 66, 131);"> Args:</span></div><div><span style="color: rgb(10, 66, 131);"> A_seq (seqaij): local submatrix of A</span></div><div><span style="color: rgb(10, 66, 131);"> B_seq (seqaij): sequential matrix B</span></div><br><div><span style="color: rgb(10, 66, 131);"> Returns:</span></div><div><span style="color: rgb(10, 66, 131);"> seq mat: PETSc matrix that is the product of A_seq and B_seq</span></div><div><span style="color: rgb(10, 66, 131);"> """</span></div><div> _, A_seq_cols <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">A_seq</span>.getSize()</div><div> B_seq_rows, _ <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">B_seq</span>.getSize()</div><div> <span style="color: rgb(168, 3, 20);">assert</span> (</div><div> A_seq_cols <span style="color: rgb(168, 3, 20);">==</span> B_seq_rows</div><div> ), <span style="color: rgb(168, 3, 20);">f</span><span style="color: rgb(10, 66, 131);">"Incompatible matrix sizes for multiplication: </span><span style="color: rgb(13, 92, 181);">{</span>A_seq_cols<span style="color: rgb(13, 92, 181);">}</span><span style="color: rgb(10, 66, 131);"> != </span><span style="color: rgb(13, 92, 181);">{</span>B_seq_rows<span style="color: rgb(13, 92, 181);">}</span><span style="color: rgb(10, 66, 131);">"</span></div><div> C_local <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">A_seq</span>.matMult(<span style="color: rgb(166, 61, 0);">B_seq</span>)</div><div> <span style="color: rgb(168, 3, 20);">return</span> C_local</div><br><br><div><span style="color: rgb(168, 3, 20);">def</span> <span style="color: rgb(72, 31, 148);">create_global_matrix</span>(<span style="color: rgb(166, 61, 0);">C_local</span>, <span style="color: rgb(166, 61, 0);">A</span>):</div><div> <span style="color: rgb(10, 66, 131);">"""Create the global matrix C from the local submatrix C_local</span></div><br><div><span style="color: rgb(10, 66, 131);"> Args:</span></div><div><span style="color: rgb(10, 66, 131);"> C_local (seqaij): local submatrix of C</span></div><div><span style="color: rgb(10, 66, 131);"> A (mpi PETSc mat): PETSc matrix A</span></div><br><div><span style="color: rgb(10, 66, 131);"> Returns:</span></div><div><span style="color: rgb(10, 66, 131);"> mpi PETSc mat: partitioned PETSc matrix C</span></div><div><span style="color: rgb(10, 66, 131);"> """</span></div><div> C_local_rows, C_local_cols <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">C_local</span>.getSize()</div><div> local_rows_start, _ <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">A</span>.getOwnershipRange()</div><div> m, _ <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">A</span>.getSize()</div><div> <span style="color: rgb(13, 92, 181);">C</span> <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.Mat().createAIJ(</div><div> <span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>((<span style="color: rgb(13, 92, 181);">None</span>, m), (C_local_cols, C_local_cols)), <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">COMM_WORLD</span></div><div> )</div><div> <span style="color: rgb(13, 92, 181);">C</span>.setUp()</div><div> <span style="color: rgb(168, 3, 20);">for</span> i <span style="color: rgb(168, 3, 20);">in</span> <span style="color: rgb(72, 31, 148);">range</span>(C_local_rows):</div><div> cols, values <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">C_local</span>.getRow(i)</div><div> global_row <span style="color: rgb(168, 3, 20);">=</span> i <span style="color: rgb(168, 3, 20);">+</span> local_rows_start</div><div> <span style="color: rgb(13, 92, 181);">C</span>.setValues(global_row, cols, values)</div><div> <span style="color: rgb(13, 92, 181);">C</span>.assemblyBegin()</div><div> <span style="color: rgb(13, 92, 181);">C</span>.assemblyEnd()</div><div> <span style="color: rgb(168, 3, 20);">return</span> <span style="color: rgb(13, 92, 181);">C</span></div><br><br><div>m, k <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">11</span>, <span style="color: rgb(13, 92, 181);">7</span></div><div><span style="color: rgb(164, 176, 188);"># Generate the random numpy matrices</span></div><div><span style="color: rgb(72, 31, 148);">np</span>.<span style="color: rgb(72, 31, 148);">random</span>.seed(<span style="color: rgb(13, 92, 181);">0</span>) <span style="color: rgb(164, 176, 188);"># sets the seed to 0</span></div><div>A_np <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">np</span>.<span style="color: rgb(72, 31, 148);">random</span>.randint(<span style="color: rgb(166, 61, 0);">low</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">0</span>, <span style="color: rgb(166, 61, 0);">high</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">6</span>, <span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>(m, k))</div><div>B_np <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">np</span>.<span style="color: rgb(72, 31, 148);">random</span>.randint(<span style="color: rgb(166, 61, 0);">low</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">0</span>, <span style="color: rgb(166, 61, 0);">high</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">6</span>, <span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>(k, k))</div><br><div><span style="color: rgb(164, 176, 188);"># Create B as a sequential matrix on each process</span></div><div>B_seq <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">create_petsc_matrix_seq</span>(B_np)</div><div><span style="color: rgb(72, 31, 148);">print_mat_info</span>(B_seq, <span style="color: rgb(10, 66, 131);">"B"</span>)</div><br><div><span style="color: rgb(13, 92, 181);">A</span> <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">create_petsc_matrix</span>(A_np)</div><div><span style="color: rgb(72, 31, 148);">print_mat_info</span>(<span style="color: rgb(13, 92, 181);">A</span>, <span style="color: rgb(10, 66, 131);">"A"</span>)</div><br><div><span style="color: rgb(164, 176, 188);"># Getting the correct local submatrix to be multiplied by B_seq</span></div><div>A_local <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">get_local_submatrix</span>(<span style="color: rgb(13, 92, 181);">A</span>)</div><br><div><span style="color: rgb(164, 176, 188);"># Multiplication of 2 sequential matrices</span></div><div>C_local <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">multiply_sequential_matrices</span>(A_local, B_seq)</div><br><div><span style="color: rgb(164, 176, 188);"># Creating the global C matrix</span></div><div><span style="color: rgb(13, 92, 181);">C</span> <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">create_global_matrix</span>(C_local, <span style="color: rgb(13, 92, 181);">A</span>)</div><div><span style="color: rgb(72, 31, 148);">print_mat_info</span>(<span style="color: rgb(13, 92, 181);">C</span>, <span style="color: rgb(10, 66, 131);">"C"</span>)</div><br><div><span style="color: rgb(164, 176, 188);"># --------------------------------------------</span></div><div><span style="color: rgb(164, 176, 188);"># TEST: Multiplication of 2 numpy matrices</span></div><div><span style="color: rgb(164, 176, 188);"># --------------------------------------------</span></div><div>AB_np <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">np</span>.<span style="color: rgb(72, 31, 148);">dot</span>(A_np, B_np)</div><div><span style="color: rgb(72, 31, 148);">Print</span>(<span style="color: rgb(168, 3, 20);">f</span><span style="color: rgb(10, 66, 131);">"MATRIX AB_np [</span><span style="color: rgb(13, 92, 181);">{</span>AB_np.shape[<span style="color: rgb(13, 92, 181);">0</span>]<span style="color: rgb(13, 92, 181);">}</span><span style="color: rgb(10, 66, 131);">x</span><span style="color: rgb(13, 92, 181);">{</span>AB_np.shape[<span style="color: rgb(13, 92, 181);">1</span>]<span style="color: rgb(13, 92, 181);">}</span><span style="color: rgb(10, 66, 131);">]"</span>)</div><div><span style="color: rgb(72, 31, 148);">Print</span>(AB_np)</div><br><div><span style="color: rgb(164, 176, 188);"># Get the local values from C</span></div><div>local_rows_start, local_rows_end <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">C</span>.getOwnershipRange()</div><div>C_local <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">C</span>.getValues(<span style="color: rgb(72, 31, 148);">range</span>(local_rows_start, local_rows_end), <span style="color: rgb(72, 31, 148);">range</span>(k))</div><br><div><span style="color: rgb(164, 176, 188);"># Assert the correctness of the multiplication for the local subset</span></div><div><span style="color: rgb(72, 31, 148);">assert_array_almost_equal</span>(C_local, AB_np[local_rows_start:local_rows_end, :], <span style="color: rgb(166, 61, 0);">decimal</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">5</span>)</div><br></div></div><div><br></div><div><br><div><br><blockquote type="cite"><div>On 23 Aug 2023, at 14:40, Pierre Jolivet <pierre.jolivet@lip6.fr> wrote:</div><br class="Apple-interchange-newline"><div><meta http-equiv="content-type" content="text/html; charset=utf-8"><div style="overflow-wrap: break-word; -webkit-nbsp-mode: space; line-break: after-white-space;"><br><div><br><blockquote type="cite"><div>On 23 Aug 2023, at 8:59 PM, Thanasis Boutsikakis <thanasis.boutsikakis@corintis.com> wrote:</div><br class="Apple-interchange-newline"><div><meta http-equiv="content-type" content="text/html; charset=utf-8"><div style="overflow-wrap: break-word; -webkit-nbsp-mode: space; line-break: after-white-space;"><div dir="auto" style="overflow-wrap: break-word; -webkit-nbsp-mode: space; line-break: after-white-space;">Thanks for the clarification Mark.<div><br></div><div>I have tried such an implementation but I since I could not find the equivalent of <a href="https://petsc.org/main/manualpages/Mat/MatMPIAIJGetLocalMat/">https://petsc.org/main/manualpages/Mat/MatMPIAIJGetLocalMat/</a> for petsc4py, I used <span style="font-family: Menlo, Monaco, "Courier New", monospace; white-space: pre; color: rgb(13, 92, 181);">A</span><span style="caret-color: rgb(29, 33, 38); color: rgb(29, 33, 38); font-family: Menlo, Monaco, "Courier New", monospace; white-space: pre; background-color: rgb(255, 255, 255);">.getLocalSubMatrix </span><font>to do so, which returns a ‘localref’ object that I cannot then use to get my local 'seqaij’ matrix.</font></div></div></div></div></blockquote><div><br></div><div>You really need MatMPIAIJGetLocalMat().</div><div>It seems there is a missing petsc4py binding, either add it yourself (and please create a merge request), or post an issue on GitLab and hope that someone adds it.</div><div>Another way to bypass the issue would be to call MatCreateSubMatrices() with an is_row set to the local rows, and an is_col set to all columns, but this will be less scalable than MatMPIAIJGetLocalMat().</div><div><br></div><div>Thanks,</div><div>Pierre</div><br><blockquote type="cite"><div><div style="overflow-wrap: break-word; -webkit-nbsp-mode: space; line-break: after-white-space;"><div dir="auto" style="overflow-wrap: break-word; -webkit-nbsp-mode: space; line-break: after-white-space;"><div><font> I also tried A.getDenseLocalMatrix() but it seems not to be suitable for my case. I couldn’t find an example in the petsc4py source code or something relevant, do you have any ideas?</font></div><div><br></div><div><div style="color: rgb(29, 33, 38); background-color: rgb(255, 255, 255); font-family: Menlo, Monaco, "Courier New", monospace; line-height: 18px; white-space: pre;"><div><div style="line-height: 18px;"><div><span style="color: rgb(10, 66, 131);">"""Experimenting with petsc mat-mat multiplication"""</span></div><div><span style="color: rgb(164, 176, 188);"># import pdb</span></div><br><div><span style="color: rgb(168, 3, 20);">import</span> <span style="color: rgb(72, 31, 148);">numpy</span> <span style="color: rgb(168, 3, 20);">as</span> <span style="color: rgb(72, 31, 148);">np</span></div><div><span style="color: rgb(168, 3, 20);">from</span> <span style="color: rgb(72, 31, 148);">firedrake</span> <span style="color: rgb(168, 3, 20);">import</span> <span style="color: rgb(13, 92, 181);">COMM_WORLD</span></div><div><span style="color: rgb(168, 3, 20);">from</span> <span style="color: rgb(72, 31, 148);">firedrake</span>.<span style="color: rgb(72, 31, 148);">petsc</span> <span style="color: rgb(168, 3, 20);">import</span> <span style="color: rgb(72, 31, 148);">PETSc</span></div><div><span style="color: rgb(168, 3, 20);">from</span> <span style="color: rgb(72, 31, 148);">numpy</span>.<span style="color: rgb(72, 31, 148);">testing</span> <span style="color: rgb(168, 3, 20);">import</span> <span style="color: rgb(72, 31, 148);">assert_array_almost_equal</span></div><div><span style="color: rgb(168, 3, 20);">import</span> <span style="color: rgb(72, 31, 148);">pdb</span></div><br><div>nproc <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">COMM_WORLD</span>.size</div><div>rank <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">COMM_WORLD</span>.rank</div><br><br><div><span style="color: rgb(168, 3, 20);">def</span> <span style="color: rgb(72, 31, 148);">Print</span>(<span style="color: rgb(166, 61, 0);">x</span>: <span style="color: rgb(72, 31, 148);">str</span>):</div><div> <span style="color: rgb(10, 66, 131);">"""Prints the string only on the root process</span></div><br><div><span style="color: rgb(10, 66, 131);"> Args:</span></div><div><span style="color: rgb(10, 66, 131);"> x (str): String to be printed</span></div><div><span style="color: rgb(10, 66, 131);"> """</span></div><div> <span style="color: rgb(72, 31, 148);">PETSc</span>.Sys.Print(<span style="color: rgb(166, 61, 0);">x</span>)</div><br><br><div><span style="color: rgb(168, 3, 20);">def</span> <span style="color: rgb(72, 31, 148);">create_petsc_matrix_seq</span>(<span style="color: rgb(166, 61, 0);">input_array</span>):</div><div> <span style="color: rgb(10, 66, 131);">"""Building a sequential petsc matrix from an array</span></div><br><div><span style="color: rgb(10, 66, 131);"> Args:</span></div><div><span style="color: rgb(10, 66, 131);"> input_array (np array): Input array</span></div><br><div><span style="color: rgb(10, 66, 131);"> Returns:</span></div><div><span style="color: rgb(10, 66, 131);"> seq mat: PETSc matrix</span></div><div><span style="color: rgb(10, 66, 131);"> """</span></div><div> <span style="color: rgb(168, 3, 20);">assert</span> <span style="color: rgb(72, 31, 148);">len</span>(<span style="color: rgb(166, 61, 0);">input_array</span>.shape) <span style="color: rgb(168, 3, 20);">==</span> <span style="color: rgb(13, 92, 181);">2</span></div><br><div> m, n <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">input_array</span>.shape</div><div> matrix <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.Mat().createAIJ(<span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>(m, n), <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(72, 31, 148);">PETSc</span>.<span style="color: rgb(13, 92, 181);">COMM_SELF</span>)</div><div> matrix.setUp()</div><br><div> matrix.setValues(<span style="color: rgb(72, 31, 148);">range</span>(m), <span style="color: rgb(72, 31, 148);">range</span>(n), <span style="color: rgb(166, 61, 0);">input_array</span>, <span style="color: rgb(166, 61, 0);">addv</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">False</span>)</div><br><div> <span style="color: rgb(164, 176, 188);"># Assembly the matrix to compute the final structure</span></div><div> matrix.assemblyBegin()</div><div> matrix.assemblyEnd()</div><br><div> <span style="color: rgb(168, 3, 20);">return</span> matrix</div><br><br><div><span style="color: rgb(168, 3, 20);">def</span> <span style="color: rgb(72, 31, 148);">create_petsc_matrix</span>(<span style="color: rgb(166, 61, 0);">input_array</span>, <span style="color: rgb(166, 61, 0);">partition_like</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">None</span>):</div><div> <span style="color: rgb(10, 66, 131);">"""Create a PETSc matrix from an input_array</span></div><br><div><span style="color: rgb(10, 66, 131);"> Args:</span></div><div><span style="color: rgb(10, 66, 131);"> input_array (np array): Input array</span></div><div><span style="color: rgb(10, 66, 131);"> partition_like (petsc mat, optional): Petsc matrix. Defaults to None.</span></div><div><span style="color: rgb(10, 66, 131);"> sparse (bool, optional): Toggle for sparese or dense. Defaults to True.</span></div><br><div><span style="color: rgb(10, 66, 131);"> Returns:</span></div><div><span style="color: rgb(10, 66, 131);"> petsc mat: PETSc matrix</span></div><div><span style="color: rgb(10, 66, 131);"> """</span></div><div> <span style="color: rgb(164, 176, 188);"># Check if input_array is 1D and reshape if necessary</span></div><div> <span style="color: rgb(168, 3, 20);">assert</span> <span style="color: rgb(72, 31, 148);">len</span>(<span style="color: rgb(166, 61, 0);">input_array</span>.shape) <span style="color: rgb(168, 3, 20);">==</span> <span style="color: rgb(13, 92, 181);">2</span>, <span style="color: rgb(10, 66, 131);">"Input array should be 2-dimensional"</span></div><div> global_rows, global_cols <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">input_array</span>.shape</div><br><div> comm <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">COMM_WORLD</span></div><div> <span style="color: rgb(168, 3, 20);">if</span> <span style="color: rgb(166, 61, 0);">partition_like</span> <span style="color: rgb(168, 3, 20);">is</span> <span style="color: rgb(168, 3, 20);">not</span> <span style="color: rgb(13, 92, 181);">None</span>:</div><div> local_rows_start, local_rows_end <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(166, 61, 0);">partition_like</span>.getOwnershipRange()</div><div> local_rows <span style="color: rgb(168, 3, 20);">=</span> local_rows_end <span style="color: rgb(168, 3, 20);">-</span> local_rows_start</div><br><div> <span style="color: rgb(164, 176, 188);"># No parallelization in the columns, set local_cols = None to parallelize</span></div><div> size <span style="color: rgb(168, 3, 20);">=</span> ((local_rows, global_rows), (global_cols, global_cols))</div><div> <span style="color: rgb(168, 3, 20);">else</span>:</div><div> size <span style="color: rgb(168, 3, 20);">=</span> ((<span style="color: rgb(13, 92, 181);">None</span>, global_rows), (global_cols, global_cols))</div><br><div> matrix <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.Mat().createAIJ(<span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>size, <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span>comm)</div><div> matrix.setUp()</div><br><div> local_rows_start, local_rows_end <span style="color: rgb(168, 3, 20);">=</span> matrix.getOwnershipRange()</div><br><div> <span style="color: rgb(168, 3, 20);">for</span> counter, i <span style="color: rgb(168, 3, 20);">in</span> <span style="color: rgb(72, 31, 148);">enumerate</span>(<span style="color: rgb(72, 31, 148);">range</span>(local_rows_start, local_rows_end)):</div><div> <span style="color: rgb(164, 176, 188);"># Calculate the correct row in the array for the current process</span></div><div> row_in_array <span style="color: rgb(168, 3, 20);">=</span> counter <span style="color: rgb(168, 3, 20);">+</span> local_rows_start</div><div> matrix.setValues(</div><div> i, <span style="color: rgb(72, 31, 148);">range</span>(global_cols), <span style="color: rgb(166, 61, 0);">input_array</span>[row_in_array, :], <span style="color: rgb(166, 61, 0);">addv</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">False</span></div><div> )</div><br><div> <span style="color: rgb(164, 176, 188);"># Assembly the matrix to compute the final structure</span></div><div> matrix.assemblyBegin()</div><div> matrix.assemblyEnd()</div><br><div> <span style="color: rgb(168, 3, 20);">return</span> matrix</div><br><br><div>m, k <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">10</span>, <span style="color: rgb(13, 92, 181);">3</span></div><div><span style="color: rgb(164, 176, 188);"># Generate the random numpy matrices</span></div><div><span style="color: rgb(72, 31, 148);">np</span>.<span style="color: rgb(72, 31, 148);">random</span>.seed(<span style="color: rgb(13, 92, 181);">0</span>) <span style="color: rgb(164, 176, 188);"># sets the seed to 0</span></div><div>A_np <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">np</span>.<span style="color: rgb(72, 31, 148);">random</span>.randint(<span style="color: rgb(166, 61, 0);">low</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">0</span>, <span style="color: rgb(166, 61, 0);">high</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">6</span>, <span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>(m, k))</div><div>B_np <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">np</span>.<span style="color: rgb(72, 31, 148);">random</span>.randint(<span style="color: rgb(166, 61, 0);">low</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">0</span>, <span style="color: rgb(166, 61, 0);">high</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">6</span>, <span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>(k, k))</div><br><div><span style="color: rgb(164, 176, 188);"># Create B as a sequential matrix on each process</span></div><div>B_seq <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">create_petsc_matrix_seq</span>(B_np)</div><br><div><span style="color: rgb(72, 31, 148);">Print</span>(B_seq.getType())</div><div><span style="color: rgb(72, 31, 148);">Print</span>(B_seq.getSizes())</div><br><div><span style="color: rgb(13, 92, 181);">A</span> <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">create_petsc_matrix</span>(A_np)</div><br><div><span style="color: rgb(72, 31, 148);">print</span>(<span style="color: rgb(10, 66, 131);">"A type:"</span>, <span style="color: rgb(13, 92, 181);">A</span>.getType())</div><div><span style="color: rgb(72, 31, 148);">print</span>(<span style="color: rgb(10, 66, 131);">"A sizes:"</span>, <span style="color: rgb(13, 92, 181);">A</span>.getSizes())</div><div><span style="color: rgb(72, 31, 148);">print</span>(<span style="color: rgb(10, 66, 131);">"A local ownership range:"</span>, <span style="color: rgb(13, 92, 181);">A</span>.getOwnershipRange())</div><br><div><span style="color: rgb(164, 176, 188);"># pdb.set_trace()</span></div><br><div><span style="color: rgb(164, 176, 188);"># Create a local sequential matrix for A using the local submatrix</span></div><div>local_rows_start, local_rows_end <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">A</span>.getOwnershipRange()</div><div>local_rows <span style="color: rgb(168, 3, 20);">=</span> local_rows_end <span style="color: rgb(168, 3, 20);">-</span> local_rows_start</div><br><div><span style="color: rgb(72, 31, 148);">print</span>(<span style="color: rgb(10, 66, 131);">"local_rows_start:"</span>, local_rows_start)</div><div><span style="color: rgb(72, 31, 148);">print</span>(<span style="color: rgb(10, 66, 131);">"local_rows_end:"</span>, local_rows_end)</div><div><span style="color: rgb(72, 31, 148);">print</span>(<span style="color: rgb(10, 66, 131);">"local_rows:"</span>, local_rows)</div><br><div>local_A <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.Mat().createAIJ(<span style="color: rgb(166, 61, 0);">size</span><span style="color: rgb(168, 3, 20);">=</span>(local_rows, k), <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(72, 31, 148);">PETSc</span>.<span style="color: rgb(13, 92, 181);">COMM_SELF</span>)</div><br><div><span style="color: rgb(164, 176, 188);"># pdb.set_trace()</span></div><br><div>comm <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">A</span>.getComm()</div><div>rows <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.IS().createStride(local_rows, <span style="color: rgb(166, 61, 0);">first</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">0</span>, <span style="color: rgb(166, 61, 0);">step</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">1</span>, <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span>comm)</div><div>cols <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.IS().createStride(k, <span style="color: rgb(166, 61, 0);">first</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">0</span>, <span style="color: rgb(166, 61, 0);">step</span><span style="color: rgb(168, 3, 20);">=</span><span style="color: rgb(13, 92, 181);">1</span>, <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span>comm)</div><br><div><span style="color: rgb(72, 31, 148);">print</span>(<span style="color: rgb(10, 66, 131);">"rows indices:"</span>, rows.getIndices())</div><div><span style="color: rgb(72, 31, 148);">print</span>(<span style="color: rgb(10, 66, 131);">"cols indices:"</span>, cols.getIndices())</div><br><div><span style="color: rgb(164, 176, 188);"># pdb.set_trace()</span></div><br><div><span style="color: rgb(164, 176, 188);"># Create the local to global mapping for rows and columns</span></div><div>l2g_rows <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.LGMap().create(rows.getIndices(), <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span>comm)</div><div>l2g_cols <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(72, 31, 148);">PETSc</span>.LGMap().create(cols.getIndices(), <span style="color: rgb(166, 61, 0);">comm</span><span style="color: rgb(168, 3, 20);">=</span>comm)</div><br><div><span style="color: rgb(72, 31, 148);">print</span>(<span style="color: rgb(10, 66, 131);">"l2g_rows type:"</span>, <span style="color: rgb(72, 31, 148);">type</span>(l2g_rows))</div><div><span style="color: rgb(72, 31, 148);">print</span>(<span style="color: rgb(10, 66, 131);">"l2g_rows:"</span>, l2g_rows.view())</div><div><span style="color: rgb(72, 31, 148);">print</span>(<span style="color: rgb(10, 66, 131);">"l2g_rows type:"</span>, <span style="color: rgb(72, 31, 148);">type</span>(l2g_cols))</div><div><span style="color: rgb(72, 31, 148);">print</span>(<span style="color: rgb(10, 66, 131);">"l2g_cols:"</span>, l2g_cols.view())</div><br><div><span style="color: rgb(164, 176, 188);"># pdb.set_trace()</span></div><br><div><span style="color: rgb(164, 176, 188);"># Set the local-to-global mapping for the matrix</span></div><div><span style="color: rgb(13, 92, 181);">A</span>.setLGMap(l2g_rows, l2g_cols)</div><br><div><span style="color: rgb(164, 176, 188);"># pdb.set_trace()</span></div><br><div><span style="color: rgb(164, 176, 188);"># Now you can get the local submatrix</span></div><div>local_A <span style="color: rgb(168, 3, 20);">=</span> <span style="color: rgb(13, 92, 181);">A</span>.getLocalSubMatrix(rows, cols)</div><br><div><span style="color: rgb(164, 176, 188);"># Assembly the matrix to compute the final structure</span></div><div>local_A.assemblyBegin()</div><div>local_A.assemblyEnd()</div><br><div><span style="color: rgb(72, 31, 148);">Print</span>(local_A.getType())</div><div><span style="color: rgb(72, 31, 148);">Print</span>(local_A.getSizes())</div><br><div><span style="color: rgb(164, 176, 188);"># pdb.set_trace()</span></div><br><div><span style="color: rgb(164, 176, 188);"># Multiply the two matrices</span></div><div>local_C <span style="color: rgb(168, 3, 20);">=</span> local_A.matMult(B_seq)</div></div></div><br></div><div><br><blockquote type="cite"><div>On 23 Aug 2023, at 12:56, Mark Adams <mfadams@lbl.gov> wrote:</div><br class="Apple-interchange-newline"><div><div dir="ltr"><div dir="ltr"><br></div><br><div class="gmail_quote"><div dir="ltr" class="gmail_attr">On Wed, Aug 23, 2023 at 5:36 AM Thanasis Boutsikakis <<a href="mailto:thanasis.boutsikakis@corintis.com">thanasis.boutsikakis@corintis.com</a>> wrote:<br></div><blockquote class="gmail_quote" style="margin:0px 0px 0px 0.8ex;border-left:1px solid rgb(204,204,204);padding-left:1ex"><div><span style="">Thanks for the suggestion Pierre.</span><div><br></div><div>Yes B is duplicated by all processes.<div><br></div><div>In this case, should B be created as a sequential sparse matrix using COMM_SELF? </div></div></div></blockquote><div><br></div><div>Yes, that is what Pierre said,</div><div><br></div><div>Mark</div><div> </div><blockquote class="gmail_quote" style="margin:0px 0px 0px 0.8ex;border-left:1px solid rgb(204,204,204);padding-left:1ex"><div><div><div>I guess if not, the multiplication of B with the output of <a href="https://petsc.org/main/manualpages/Mat/MatMPIAIJGetLocalMat/" target="_blank">https://petsc.org/main/manualpages/Mat/MatMPIAIJGetLocalMat/</a> would not go through, right? </div><div><br></div><div>Thanks,</div><div>Thanos</div><div><div> <br><blockquote type="cite"><div>On 23 Aug 2023, at 10:47, Pierre Jolivet <<a href="mailto:pierre.jolivet@lip6.fr" target="_blank">pierre.jolivet@lip6.fr</a>> wrote:</div><br><div><div dir="auto"><div dir="ltr"><div dir="ltr"></div><div dir="ltr"><br></div><div dir="ltr"><blockquote type="cite">On 23 Aug 2023, at 5:35 PM, Thanasis Boutsikakis <<a href="mailto:thanasis.boutsikakis@corintis.com" target="_blank">thanasis.boutsikakis@corintis.com</a>> wrote:<br><br></blockquote></div><blockquote type="cite"><div dir="ltr">Hi all,<div><br></div><div>I am trying to multiply two Petsc matrices as C = A * B, where A is a tall matrix and B is a relatively small matrix.</div><div><br></div><div>I have taken the decision to create A as (row-)partitioned matrix and B as a non-partitioned matrix that it is entirely shared by all procs (to avoid unnecessary communication).</div><div><br></div><div>Here is my code:</div><div><br></div><div style="color:rgb(29,33,38);background-color:rgb(255,255,255);font-family:Menlo,Monaco,"Courier New",monospace;line-height:18px;white-space:pre-wrap"><div><span style="color:rgb(168,3,20)">import</span> <span style="color:rgb(72,31,148)">numpy</span> <span style="color:rgb(168,3,20)">as</span> <span style="color:rgb(72,31,148)">np</span></div><div><span style="color:rgb(168,3,20)">from</span> <span style="color:rgb(72,31,148)">firedrake</span> <span style="color:rgb(168,3,20)">import</span> <span style="color:rgb(13,92,181)">COMM_WORLD</span></div><div><span style="color:rgb(168,3,20)">from</span> <span style="color:rgb(72,31,148)">firedrake</span>.<span style="color:rgb(72,31,148)">petsc</span> <span style="color:rgb(168,3,20)">import</span> <span style="color:rgb(72,31,148)">PETSc</span></div><div><span style="color:rgb(168,3,20)">from</span> <span style="color:rgb(72,31,148)">numpy</span>.<span style="color:rgb(72,31,148)">testing</span> <span style="color:rgb(168,3,20)">import</span> <span style="color:rgb(72,31,148)">assert_array_almost_equal</span></div><br><div>nproc <span style="color:rgb(168,3,20)">=</span> <span style="color:rgb(13,92,181)">COMM_WORLD</span>.size</div><div>rank <span style="color:rgb(168,3,20)">=</span> <span style="color:rgb(13,92,181)">COMM_WORLD</span>.rank</div><br><div><span style="color:rgb(168,3,20)">def</span> <span style="color:rgb(72,31,148)">create_petsc_matrix_non_partitioned</span>(<span style="color:rgb(166,61,0)">input_array</span>):</div><div> <span style="color:rgb(10,66,131)">"""Building a mpi non-partitioned petsc matrix from an array</span></div><br><div><span style="color:rgb(10,66,131)"> Args:</span></div><div><span style="color:rgb(10,66,131)"> input_array (np array): Input array</span></div><div><span style="color:rgb(10,66,131)"> sparse (bool, optional): Toggle for sparese or dense. Defaults to True.</span></div><br><div><span style="color:rgb(10,66,131)"> Returns:</span></div><div><span style="color:rgb(10,66,131)"> mpi mat: PETSc matrix</span></div><div><span style="color:rgb(10,66,131)"> """</span></div><div> <span style="color:rgb(168,3,20)">assert</span> <span style="color:rgb(72,31,148)">len</span>(<span style="color:rgb(166,61,0)">input_array</span>.shape) <span style="color:rgb(168,3,20)">==</span> <span style="color:rgb(13,92,181)">2</span></div><br><div> m, n <span style="color:rgb(168,3,20)">=</span> <span style="color:rgb(166,61,0)">input_array</span>.shape</div><br><div> matrix <span style="color:rgb(168,3,20)">=</span> <span style="color:rgb(72,31,148)">PETSc</span>.Mat().createAIJ(<span style="color:rgb(166,61,0)">size</span><span style="color:rgb(168,3,20)">=</span>((m, n), (m, n)), <span style="color:rgb(166,61,0)">comm</span><span style="color:rgb(168,3,20)">=</span><span style="color:rgb(13,92,181)">COMM_WORLD</span>)</div><br><div> <span style="color:rgb(164,176,188)"># Set the values of the matrix</span></div><div> matrix.setValues(<span style="color:rgb(72,31,148)">range</span>(m), <span style="color:rgb(72,31,148)">range</span>(n), <span style="color:rgb(166,61,0)">input_array</span>[:, :], <span style="color:rgb(166,61,0)">addv</span><span style="color:rgb(168,3,20)">=</span><span style="color:rgb(13,92,181)">False</span>)</div><br><div> <span style="color:rgb(164,176,188)"># Assembly the matrix to compute the final structure</span></div><div> matrix.assemblyBegin()</div><div> matrix.assemblyEnd()</div><br><div> <span style="color:rgb(168,3,20)">return</span> matrix</div><br><br><div><span style="color:rgb(168,3,20)">def</span> <span style="color:rgb(72,31,148)">create_petsc_matrix</span>(<span style="color:rgb(166,61,0)">input_array</span>, <span style="color:rgb(166,61,0)">partition_like</span><span style="color:rgb(168,3,20)">=</span><span style="color:rgb(13,92,181)">None</span>):</div><div> <span style="color:rgb(10,66,131)">"""Create a PETSc matrix from an input_array</span></div><br><div><span style="color:rgb(10,66,131)"> Args:</span></div><div><span style="color:rgb(10,66,131)"> input_array (np array): Input array</span></div><div><span style="color:rgb(10,66,131)"> partition_like (petsc mat, optional): Petsc matrix. Defaults to None.</span></div><div><span style="color:rgb(10,66,131)"> sparse (bool, optional): Toggle for sparese or dense. Defaults to True.</span></div><br><div><span style="color:rgb(10,66,131)"> Returns:</span></div><div><span style="color:rgb(10,66,131)"> petsc mat: PETSc matrix</span></div><div><span style="color:rgb(10,66,131)"> """</span></div><div> <span style="color:rgb(164,176,188)"># Check if input_array is 1D and reshape if necessary</span></div><div> <span style="color:rgb(168,3,20)">assert</span> <span style="color:rgb(72,31,148)">len</span>(<span style="color:rgb(166,61,0)">input_array</span>.shape) <span style="color:rgb(168,3,20)">==</span> <span style="color:rgb(13,92,181)">2</span>, <span style="color:rgb(10,66,131)">"Input array should be 2-dimensional"</span></div><div> global_rows, global_cols <span style="color:rgb(168,3,20)">=</span> <span style="color:rgb(166,61,0)">input_array</span>.shape</div><br><div> comm <span style="color:rgb(168,3,20)">=</span> <span style="color:rgb(13,92,181)">COMM_WORLD</span></div><div> <span style="color:rgb(168,3,20)">if</span> <span style="color:rgb(166,61,0)">partition_like</span> <span style="color:rgb(168,3,20)">is</span> <span style="color:rgb(168,3,20)">not</span> <span style="color:rgb(13,92,181)">None</span>:</div><div> local_rows_start, local_rows_end <span style="color:rgb(168,3,20)">=</span> <span style="color:rgb(166,61,0)">partition_like</span>.getOwnershipRange()</div><div> local_rows <span style="color:rgb(168,3,20)">=</span> local_rows_end <span style="color:rgb(168,3,20)">-</span> local_rows_start</div><br><div> <span style="color:rgb(164,176,188)"># No parallelization in the columns, set local_cols = None to parallelize</span></div><div> size <span style="color:rgb(168,3,20)">=</span> ((local_rows, global_rows), (global_cols, global_cols))</div><div> <span style="color:rgb(168,3,20)">else</span>:</div><div> size <span style="color:rgb(168,3,20)">=</span> ((<span style="color:rgb(13,92,181)">None</span>, global_rows), (global_cols, global_cols))</div><br><div> matrix <span style="color:rgb(168,3,20)">=</span> <span style="color:rgb(72,31,148)">PETSc</span>.Mat().createAIJ(<span style="color:rgb(166,61,0)">size</span><span style="color:rgb(168,3,20)">=</span>size, <span style="color:rgb(166,61,0)">comm</span><span style="color:rgb(168,3,20)">=</span>comm)</div><div> matrix.setUp()</div><br><div> local_rows_start, local_rows_end <span style="color:rgb(168,3,20)">=</span> matrix.getOwnershipRange()</div><br><div> <span style="color:rgb(168,3,20)">for</span> counter, i <span style="color:rgb(168,3,20)">in</span> <span style="color:rgb(72,31,148)">enumerate</span>(<span style="color:rgb(72,31,148)">range</span>(local_rows_start, local_rows_end)):</div><div> <span style="color:rgb(164,176,188)"># Calculate the correct row in the array for the current process</span></div><div> row_in_array <span style="color:rgb(168,3,20)">=</span> counter <span style="color:rgb(168,3,20)">+</span> local_rows_start</div><div> matrix.setValues(</div><div> i, <span style="color:rgb(72,31,148)">range</span>(global_cols), <span style="color:rgb(166,61,0)">input_array</span>[row_in_array, :], <span style="color:rgb(166,61,0)">addv</span><span style="color:rgb(168,3,20)">=</span><span style="color:rgb(13,92,181)">False</span></div><div> )</div><br><div> <span style="color:rgb(164,176,188)"># Assembly the matrix to compute the final structure</span></div><div> matrix.assemblyBegin()</div><div> matrix.assemblyEnd()</div><br><div> <span style="color:rgb(168,3,20)">return</span> matrix</div><br><br><div>m, k <span style="color:rgb(168,3,20)">=</span> <span style="color:rgb(13,92,181)">10</span>, <span style="color:rgb(13,92,181)">3</span></div><div><span style="color:rgb(164,176,188)"># Generate the random numpy matrices</span></div><div><span style="color:rgb(72,31,148)">np</span>.<span style="color:rgb(72,31,148)">random</span>.seed(<span style="color:rgb(13,92,181)">0</span>) <span style="color:rgb(164,176,188)"># sets the seed to 0</span></div><div>A_np <span style="color:rgb(168,3,20)">=</span> <span style="color:rgb(72,31,148)">np</span>.<span style="color:rgb(72,31,148)">random</span>.randint(<span style="color:rgb(166,61,0)">low</span><span style="color:rgb(168,3,20)">=</span><span style="color:rgb(13,92,181)">0</span>, <span style="color:rgb(166,61,0)">high</span><span style="color:rgb(168,3,20)">=</span><span style="color:rgb(13,92,181)">6</span>, <span style="color:rgb(166,61,0)">size</span><span style="color:rgb(168,3,20)">=</span>(m, k))</div><div>B_np <span style="color:rgb(168,3,20)">=</span> <span style="color:rgb(72,31,148)">np</span>.<span style="color:rgb(72,31,148)">random</span>.randint(<span style="color:rgb(166,61,0)">low</span><span style="color:rgb(168,3,20)">=</span><span style="color:rgb(13,92,181)">0</span>, <span style="color:rgb(166,61,0)">high</span><span style="color:rgb(168,3,20)">=</span><span style="color:rgb(13,92,181)">6</span>, <span style="color:rgb(166,61,0)">size</span><span style="color:rgb(168,3,20)">=</span>(k, k))</div><br><br><div><span style="color:rgb(13,92,181)">A</span> <span style="color:rgb(168,3,20)">=</span> <span style="color:rgb(72,31,148)">create_petsc_matrix</span>(A_np)</div><br><div><span style="color:rgb(13,92,181)">B</span> <span style="color:rgb(168,3,20)">=</span> <span style="color:rgb(72,31,148)">create_petsc_matrix_non_partitioned</span>(B_np)</div><br><div><span style="color:rgb(164,176,188)"># Now perform the multiplication</span></div></div><div><span style="font-family:Menlo,Monaco,"Courier New",monospace;white-space:pre-wrap;color:rgb(13,92,181)">C</span><span style="color:rgb(29,33,38);font-family:Menlo,Monaco,"Courier New",monospace;white-space:pre-wrap;background-color:rgb(255,255,255)"> </span><span style="font-family:Menlo,Monaco,"Courier New",monospace;white-space:pre-wrap;color:rgb(168,3,20)">=</span><span style="color:rgb(29,33,38);font-family:Menlo,Monaco,"Courier New",monospace;white-space:pre-wrap;background-color:rgb(255,255,255)"> </span><span style="font-family:Menlo,Monaco,"Courier New",monospace;white-space:pre-wrap;color:rgb(13,92,181)">A</span><span style="color:rgb(29,33,38);font-family:Menlo,Monaco,"Courier New",monospace;white-space:pre-wrap;background-color:rgb(255,255,255)"> </span><span style="font-family:Menlo,Monaco,"Courier New",monospace;white-space:pre-wrap;color:rgb(168,3,20)">*</span><span style="color:rgb(29,33,38);font-family:Menlo,Monaco,"Courier New",monospace;white-space:pre-wrap;background-color:rgb(255,255,255)"> </span><span style="font-family:Menlo,Monaco,"Courier New",monospace;white-space:pre-wrap;color:rgb(13,92,181)">B</span></div><div><br></div><div>The problem with this is that there is a mismatch between the local rows of A (depend on the partitioning) and the global rows of B (3 for all procs), so that the multiplication cannot happen in parallel. Here is the error:</div><div><br></div><div><div>[0]PETSC ERROR: ------------------------------------------------------------------------</div><div>[0]PETSC ERROR: Caught signal number 11 SEGV: Segmentation Violation, probably memory access out of range</div><div>[0]PETSC ERROR: Try option -start_in_debugger or -on_error_attach_debugger</div><div>[0]PETSC ERROR: or see <a href="https://petsc.org/release/faq/#valgrind" target="_blank">https://petsc.org/release/faq/#valgrind</a> and <a href="https://petsc.org/release/faq/" target="_blank">https://petsc.org/release/faq/</a></div><div>[0]PETSC ERROR: configure using --with-debugging=yes, recompile, link, and run</div><div>[0]PETSC ERROR: to get more information on the crash.</div><div>[0]PETSC ERROR: Run with -malloc_debug to check if memory corruption is causing the crash.</div><div>application called MPI_Abort(MPI_COMM_WORLD, 59) - process 0</div><div>[unset]: write_line error; fd=-1 buf=:cmd=abort exitcode=59</div><div>:</div><div>system msg for write_line failure : Bad file descriptor</div></div><div><br></div><div><br></div><div>Is there a standard way to achieve this?</div></div></blockquote><div><br></div><div>Your B is duplicated by all processes?</div><div>If so, then, call <a href="https://petsc.org/main/manualpages/Mat/MatMPIAIJGetLocalMat/" target="_blank">https://petsc.org/main/manualpages/Mat/MatMPIAIJGetLocalMat/</a>, do a sequential product with B on COMM_SELF, not COMM_WORLD, and use <a href="https://petsc.org/main/manualpages/Mat/MatCreateMPIMatConcatenateSeqMat/" target="_blank">https://petsc.org/main/manualpages/Mat/MatCreateMPIMatConcatenateSeqMat/</a> with the output.</div><div><br></div><div>Thanks,</div><div>Pierre</div><br><blockquote type="cite"><div dir="ltr"><div>Thanks,</div><div>Thanos</div><div><br></div><div> </div></div></blockquote></div></div></div></blockquote></div><br></div></div></div></blockquote></div></div>
</div></blockquote></div><br></div></div></div></div></blockquote></div><br></div></div></blockquote></div><br></div></div></body></html>