[query] Add tree_matmul, matrix multiply in case of large inner dimension #9063

johnc1231 · 2020-07-08T17:28:23Z

This PR introduces tree_matmul, a method on BlockMatrix that allows for greater parallelism when multiplying two large matrices that result in a small matrix (i.e. the inner dimension is much larger than the outer dimensions).

In order to do this, this PR takes a split_on_inner parameter. This parameter defines how many subdivisions to break the two matrices into. Corresponding subdivisions are multiplied, written to disk, then read back in and summed.

For example, if you are multiplying a 4k by 500k matrix by its transpose, your answer will be a 4k by 4k block. With normal matrix multiply, you'd get no parallelism at all, since the result is only a single partition. If you instead use tree_matmul with split_on_inner = 5, hail will do the following:

Break up 4k by 500k matrix into five 4k by 100k matrices.
Break up 500k by 4k matrix into five 100k by 4k matrices.
Multiply corresponding matrices together, creating five 4k by 4k matrices.
Write them all to disk, then read them back in.
Sum the 5 read in matrices together to get your result.

This PR also introduces a write_block_matrices function that writes out a list of BlockMatrix in parallel. This was necessary to write out all of the intermediate matrices in tree_matmul.

@konradjk
@danking

…rite many RDDs in parallel

…n files

danking

This is absolutely awesome John! Great work

danking · 2020-07-09T15:25:54Z

hail/python/test/hail/linalg/test_linalg.py

+        self._assert_eq(m.T.tree_matmul(m, 2, new_temp_file()), nm.T @ nm)
+        self._assert_eq(m.T.tree_matmul(nm, 2, new_temp_file()), nm.T @ nm)
+        self._assert_eq(row.T.tree_matmul(row, 2, new_temp_file()), nrow.T @ nrow)
+        self._assert_eq(row.T.tree_matmul(nrow, 2, new_temp_file()), nrow.T @ nrow)


call me paranoid, but it would be nice to have at least one example of x @ y where x and y are both block matrices and neither is the transpose of the other.

Heh, I just copied the current matmul tests plus that looping of different block sizes, but you're right, these are only transposes. I'll add another test.

danking · 2020-07-09T15:26:31Z

hail/src/main/scala/is/hail/linalg/BlockMatrix.scala

+                           header: Option[String],
+                           addIndex: Boolean,
+                           compression: Option[String],
+                           customFilenames: Option[Array[String]]): Unit = {


this is an odd formatting change? Maybe add a newline before the ) to get the desired formatting from IntelliJ?

Thanks, this was an accident.

danking · 2020-07-09T15:30:38Z

hail/python/hail/linalg/blockmatrix.py

@@ -1427,6 +1459,47 @@ def __matmul__(self, b):

        return BlockMatrix(BlockMatrixDot(self._bmir, b._bmir))

+    @typecheck_method(b=oneof(np.ndarray, block_matrix_type), split_on_inner=int, path_prefix=str)
+    def tree_matmul(self, b, split_on_inner, path_prefix):
+        """Matrix multiplication in situations with large inner dimension. This function splits a single matrix


We generally prefer a one sentence description on its own line, followed by a full description. Some python tools use this in tooltips

danking · 2020-07-09T15:31:43Z

hail/python/hail/linalg/blockmatrix.py

@@ -1427,6 +1459,47 @@ def __matmul__(self, b):

        return BlockMatrix(BlockMatrixDot(self._bmir, b._bmir))

+    @typecheck_method(b=oneof(np.ndarray, block_matrix_type), split_on_inner=int, path_prefix=str)
+    def tree_matmul(self, b, split_on_inner, path_prefix):


I think adding a default argument for path_prefix of None and creating a temp file in that case will improve usability.

danking · 2020-07-09T15:32:49Z

hail/python/hail/linalg/blockmatrix.py

@@ -1427,6 +1459,47 @@ def __matmul__(self, b):

        return BlockMatrix(BlockMatrixDot(self._bmir, b._bmir))

+    @typecheck_method(b=oneof(np.ndarray, block_matrix_type), split_on_inner=int, path_prefix=str)
+    def tree_matmul(self, b, split_on_inner, path_prefix):


I think we should make either the last or the last two parameters keyword-only arguments, so: (self, b, *, split_on_inner, path_prefix).

I'd be willing to make path_prefix a keyword only argument, but can you explain why you'd prefer that? I think split_on_inner I'll leave the way it is, since you always have to specify it. path_prefix with a default value of None as your above comment suggests makes the keyword only argument sound better.

d

johnc1231 · 2020-07-09T19:32:14Z

I addressed most of your comments, I'm not sure why you want the keyword only arguments though.

danking · 2020-07-09T20:53:33Z

It's kind of a gut reaction for me. There's some conversation at PEP 1302's rationale. For me, everything is about backwards compatibility and the two cases are:

changing from positional arguments without varargs to varargs and keywords is backwards-incompatible.
changing a positional-without-default to positional-with-default values is backwards-incompatible unless the positional argument is the last one.

In this concrete example, I worry about two possible evolutions of this function.

Accept multiple other matrices:

def tree_matmul(self, *others, split_on_inner, path_prefix)

Allow the compiler/python code to choose a good value for split_on_inner

def tree_matmul(self, b, split_on_inner=None, path_prefix=None) # now path prefix *must* be optional as well.

danking · 2020-07-09T20:55:13Z

I might also call split_on_inner splits for brevity in the keyword case.

patrick-schultz · 2020-07-20T16:33:31Z

@danking Should I review this too, or do you feel like you covered it?

danking · 2020-07-20T17:01:09Z

@patrick-schultz I only looked at the Python interface, so I would much appreciate your thoughts on the Scala stuff.

johnc1231 · 2020-07-21T14:21:22Z

Tagging WIP because I have two tiny changes Dan asked for to make (argument name and keyword only)

…sion (hail-is#9063) * Set up some range computations * Did most of the python work, stuck now beacuse I need to be able to write many RDDs in parallel * Kept Dan's coschedule actions thing somewhere I won't lose it * Fixed bug in copy of BlockMatrixMultiWrite * Python side BlockMatrixMultiWriter exists * Python to scala connection for block matrix native writer is working * Metadata and success files being written, just need to write partition files * Organized correctly now * Write multiple works * Split multiplying works * Fixed last_rows and last_cols computation * Added a print * Now it's tree_matmul * Began documenting the new functions * Delete coscheduleActions comment * Delete old comments * Python half of supporting stage_locally * Deleted unusued partitionURI * WIP * Updated to use using fs.create * Some typechecks fixed * Added tests for tree_matmul * Refactored, passing tests * Refactored to pass along ExecuteContext to get localTmpDir * Deleted print * pylint fixes * Restored tempfile import * Fixed indentation * Removed more accidental formatting changes * Added default path * Update the tests * Rename split_on_inner to splits * Keyword only arguments * Fix typecheck on splits

johnc1231 added 25 commits July 7, 2020 15:14

Set up some range computations

553db1b

Did most of the python work, stuck now beacuse I need to be able to w…

45eb1b6

…rite many RDDs in parallel

Kept Dan's coschedule actions thing somewhere I won't lose it

0160199

Fixed bug in copy of BlockMatrixMultiWrite

35e3de6

Python side BlockMatrixMultiWriter exists

ff93535

Python to scala connection for block matrix native writer is working

74a2fdc

Metadata and success files being written, just need to write partitio…

7dd9b03

…n files

Organized correctly now

73ee289

Write multiple works

c18f1ce

Split multiplying works

92ae540

Fixed last_rows and last_cols computation

6f06316

Added a print

0d06d3e

Now it's tree_matmul

5c25fec

Began documenting the new functions

d8b34c6

Delete coscheduleActions comment

e0c2e48

Delete old comments

4edd104

Python half of supporting stage_locally

d2a70b1

Deleted unusued partitionURI

f6b9654

WIP

7c2f5d4

Updated to use using fs.create

d6aab53

Some typechecks fixed

d56c18d

Added tests for tree_matmul

6d8fa3e

Refactored, passing tests

10d7081

Refactored to pass along ExecuteContext to get localTmpDir

8fcb56c

Deleted print

d6aa108

johnc1231 assigned patrick-schultz Jul 8, 2020

johnc1231 added 2 commits July 8, 2020 14:10

pylint fixes

1c5091f

Restored tempfile import

4151de5

danking previously requested changes Jul 9, 2020

View reviewed changes

Fixed indentation

a1edf0b

johnc1231 added 3 commits July 9, 2020 13:17

Removed more accidental formatting changes

77728cf

Added default path

f36a99d

Update the tests

bcd53ba

patrick-schultz approved these changes Jul 21, 2020

View reviewed changes

johnc1231 added the WIP label Jul 21, 2020

johnc1231 added 3 commits July 22, 2020 10:18

Rename split_on_inner to splits

92d6839

Keyword only arguments

b395efc

Fix typecheck on splits

1a55c9a

johnc1231 removed the WIP label Jul 22, 2020

danking merged commit 6360d5b into hail-is:master Jul 22, 2020

johnc1231 mentioned this pull request Jul 23, 2020

[release] Update changelog for 0.2.50. #9125

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[query] Add tree_matmul, matrix multiply in case of large inner dimension #9063

[query] Add tree_matmul, matrix multiply in case of large inner dimension #9063

johnc1231 commented Jul 8, 2020 •

edited

danking left a comment

danking Jul 9, 2020

johnc1231 Jul 9, 2020

danking Jul 9, 2020

johnc1231 Jul 9, 2020

danking Jul 9, 2020

danking Jul 9, 2020

danking Jul 9, 2020

johnc1231 Jul 9, 2020

johnc1231 commented Jul 9, 2020

danking commented Jul 9, 2020 •

edited

danking commented Jul 9, 2020

patrick-schultz commented Jul 20, 2020

danking commented Jul 20, 2020

johnc1231 commented Jul 21, 2020

[query] Add tree_matmul, matrix multiply in case of large inner dimension #9063

[query] Add tree_matmul, matrix multiply in case of large inner dimension #9063

Conversation

johnc1231 commented Jul 8, 2020 • edited

danking left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

johnc1231 commented Jul 9, 2020

danking commented Jul 9, 2020 • edited

danking commented Jul 9, 2020

patrick-schultz commented Jul 20, 2020

danking commented Jul 20, 2020

johnc1231 commented Jul 21, 2020

johnc1231 commented Jul 8, 2020 •

edited

danking commented Jul 9, 2020 •

edited