[WIP]: Jax tensor donation #3224

mamanain · 2024-09-27T17:40:47Z

This is still work in progress, just wanted to get some feedback for this version as well.

The main area of improvement: donation logic should be more in line with the jax documentation.
There are currently two problems here:

Donated buffers can still be used in the function but here it can be used as a buffer straight away. This can lead to data being overwritten and then used for some computations which will lead to errors. So we need to change the logic so that the buffer becomes available only after its last usage.
Right now buffer can be used only once. Should we keep it in the dictionary and reuse in other places where it fits?

codecov · 2024-09-27T17:46:14Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 90.03%. Comparing base (70fa878) to head (8fc4d95).
Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3224      +/-   ##
==========================================
+ Coverage   90.01%   90.03%   +0.02%     
==========================================
  Files         445      446       +1     
  Lines       55850    55953     +103     
  Branches     5351     5357       +6     
==========================================
+ Hits        50274    50378     +104     
  Misses       4169     4169              
+ Partials     1407     1406       -1

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

superlopuh

Thank you for this, it's a great first step! A few bigger comments before the nitty-gritty:

I'm a little confused by the name of the pass, it's not really converting JAX to linalg, just making sure that the operands marked as donating are actually reused as destinations, so another name feels more appropriate. (I would recommend making the prefix of the pass something that's not convert, as that's usually used for converting one dialect to another, and in this case it's more of an optimisation.) Maybe jax-use-donated-arguments?
The test is quite specific, and uses operations from dialects that aren't strictly involved. I would recommend using the "test" dialect as much as possible to generate values that are necessary for the test.
It's quite a bit more difficult to review PRs with lots of Pyright errors, could you please fix the errors locally and ping me again to take a look?
It would be good to be very clear about the limitations of the proposed approach, and to raise helpful messages if used in an unexpected context, could you please add DiagnosticExceptions if the assumptions are not met and tests for those?

superlopuh · 2024-09-29T13:52:16Z

xdsl/tools/command_line_tool.py

@@ -86,6 +86,11 @@ def get_convert_stencil_to_ll_mlir():

        return convert_stencil_to_ll_mlir.ConvertStencilToLLMLIRPass

+    def get_jax_use_donated_arguments():
+        from xdsl.transforms.experimental import jax_use_donated_arguments


I would recommend against experimental, it's a place for untested code, that we will probably remove at some point, let's just put it in the main transforms folder

superlopuh · 2024-09-29T17:27:02Z

xdsl/transforms/jax_use_donated_arguments.py

+        donated_inputs: list[BlockArgument] = []
+        for inp, attr in zip(op.args, op.arg_attrs):
+            if type(inp.type) is TensorType and "tf.aliasing_output" in attr.data:
+                donated_inputs.append(inp)


Does this work?

Suggested change

donated_inputs: list[BlockArgument] = []

for inp, attr in zip(op.args, op.arg_attrs):

if type(inp.type) is TensorType and "tf.aliasing_output" in attr.data:

donated_inputs.append(inp)

donated_inputs = [

inp

for inp, attr in zip(op.args, op.arg_attrs)

if isinstance(inp.type, TensorType) and "tf.aliasing_output" in attr.data

]

superlopuh · 2024-09-29T17:27:46Z

xdsl/transforms/jax_use_donated_arguments.py

+            walk_regions_first=True,
+        )
+        the_one_pass.rewrite_module(op)
+        MLIROptPass(arguments=("--linalg-fuse-elementwise-ops",)).apply(ctx, op)


This should be done separately by the user, and not included in this pass

superlopuh · 2024-09-29T17:28:27Z

tests/filecheck/mlir-conversion/with-mlir/jax-use-donated-arguments.mlir

there's no need to put the test in the with-mlir folder, it's better to separate it into a pure xDSL thing, and let the user call into mlir-opt if they want to

superlopuh · 2024-10-08T15:06:24Z

xdsl/transforms/jax_use_donated_arguments.py

+        arg_attrs = getattr(func_op, "arg_attrs")
+        args = getattr(func_op, "args")


we try to avoid getattr when possible, could you please check that the parent is a func, and then access properties directly?

xdsl/dialects/builtin.py

superlopuh · 2024-10-14T14:49:17Z

xdsl/dialects/tensor.py

+@irdl_op_definition
+class CollapseShapeOp(IRDLOperation):
+    name = "tensor.collapse_shape"
+
+    src = operand_def(TensorType[Attribute])
+    result = result_def(TensorType[Attribute])
+    reassociation = prop_def(ReassociationAttr)
+    assembly_format = (
+        "$src $reassociation attr-dict `:` type($src) `into` type($result)"
+    )
+
+    traits = frozenset([NoMemoryEffect()])


I don't think we need this as part of this PR

Would be great to add in a separate PR, though

xdsl/transforms/jax_use_donated_arguments.py

tests/filecheck/transforms/jax-use-donated-arguments.mlir

superlopuh · 2024-10-14T14:53:03Z

xdsl/dialects/bufferization.py

@@ -168,12 +168,33 @@ class ToMemrefOp(IRDLOperation):
    assembly_format = "$tensor (`read_only` $read_only^)?  `:` attr-dict type($memref)"


+# now only works for (tensor, tensor) arguments. need to add memref support as well.
+@irdl_op_definition
+class MaterializeInDestination(IRDLOperation):


Can you please add this in a separate PR?

superlopuh · 2024-10-14T15:37:15Z

tests/filecheck/transforms/jax-use-donated-arguments.mlir

+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+
+builtin.module {


Suggested change

builtin.module {

And in other tests

superlopuh · 2024-10-14T15:37:51Z

tests/filecheck/transforms/jax-use-donated-arguments.mlir

+// RUN: xdsl-opt %s -p jax-use-donated-arguments --split-input-file --verify-diagnostics | filecheck %s
+
+builtin.module {
+func.func public @main(%arg0: tensor<2x3xf32>, %arg1: tensor<3x4xf32>, %arg2: tensor<2x4xf32> {tf.aliasing_output = 0 : i32}) -> (tensor<2x4xf32>) {


I don't think all the functions have to be called "main", and we can use the function name to document the purpose of the test, like so:

Suggested change

func.func public @main(%arg0: tensor<2x3xf32>, %arg1: tensor<3x4xf32>, %arg2: tensor<2x4xf32> {tf.aliasing_output = 0 : i32}) -> (tensor<2x4xf32>) {

func.func public @one_arg(%arg0: tensor<2x3xf32>, %arg1: tensor<3x4xf32>, %arg2: tensor<2x4xf32> {tf.aliasing_output = 0 : i32}) -> (tensor<2x4xf32>) {

superlopuh · 2024-10-14T15:39:31Z

xdsl/dialects/bufferization.py

+    name = "bufferization.materialize_in_destination"
+
+    source = operand_def(
+        TensorMemrefInferenceConstraint("T", AnyOf([TensorType, UnrankedTensorType]))


Suggested change

TensorMemrefInferenceConstraint("T", AnyOf([TensorType, UnrankedTensorType]))

TensorMemrefInferenceConstraint("T", AnyTensorTypeConstr | AnyUnrankedTensorTypeConstr)

(will need to add the second constraint definition in builtin.py)

superlopuh · 2024-10-14T15:39:51Z

xdsl/transforms/jax_use_donated_arguments.py

+
+        donated_inputs = [
+            inp
+            for inp, attr in zip(func_op.args, func_op.arg_attrs)


Suggested change

for inp, attr in zip(func_op.args, func_op.arg_attrs)

for inp, attr in zip(func_op.args, func_op.arg_attrs, strict=True)

superlopuh · 2024-10-14T15:40:44Z

xdsl/transforms/jax_use_donated_arguments.py

+    @op_type_rewrite_pattern
+    def match_and_rewrite(self, op: Return, rewriter: PatternRewriter, /):
+        func_op = op.parent_op()
+        if func_op is None or type(func_op) is not FuncOp:


Suggested change

if func_op is None or type(func_op) is not FuncOp:

assert isinstance(func_op, FuncOp)

The input can be assumed to be verified. There's no way to handle a VerifyException here anyway, so we might as well assert.

superlopuh marked this pull request as draft September 28, 2024 22:50

superlopuh assigned mamanain Sep 28, 2024

superlopuh reviewed Sep 28, 2024

View reviewed changes

superlopuh reviewed Sep 29, 2024

View reviewed changes

superlopuh reviewed Oct 8, 2024

View reviewed changes

superlopuh reviewed Oct 14, 2024

View reviewed changes

xdsl/dialects/builtin.py Outdated Show resolved Hide resolved

superlopuh reviewed Oct 14, 2024

View reviewed changes

xdsl/transforms/jax_use_donated_arguments.py Outdated Show resolved Hide resolved

superlopuh reviewed Oct 14, 2024

View reviewed changes

xdsl/transforms/jax_use_donated_arguments.py Outdated Show resolved Hide resolved

superlopuh reviewed Oct 14, 2024

View reviewed changes

tests/filecheck/transforms/jax-use-donated-arguments.mlir Outdated Show resolved Hide resolved

superlopuh reviewed Oct 14, 2024

View reviewed changes

tests/filecheck/transforms/jax-use-donated-arguments.mlir Show resolved Hide resolved

superlopuh reviewed Oct 14, 2024

View reviewed changes

manainen added 9 commits October 15, 2024 15:21

dumb version for one operation type

dfa08fe

test added

c83999f

mlir step added

67b8a07

more general version

2ffda16

pyright fixes

121548f

change names

1319404

move test to mlir ones

6513a62

trim test

7c3d384

move file

432ebd2

manainen and others added 13 commits October 15, 2024 15:23

test mlirs

290adb3

rewrite pass to use materialize_in_destination

ad44505

move test

7fb15af

little flow fix

c83f6fa

little flow fix

c77ffda

full integration test

57a98ee

cleanup

c7279fe

little fix

a720e3c

fixed unnecessary onnx marimo changes

7459e67

style

d95168d

fixes and simplifications

9f6ae38

tests and some style

9b1bf3e

broken parsing example

7869d09

mamanain force-pushed the jax-tensor-donation branch from 93a8703 to 7869d09 Compare October 15, 2024 14:24

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[WIP]: Jax tensor donation #3224

[WIP]: Jax tensor donation #3224

mamanain commented Sep 27, 2024

codecov bot commented Sep 27, 2024 •

edited

Loading

superlopuh left a comment

superlopuh Sep 29, 2024

superlopuh Sep 29, 2024

superlopuh Sep 29, 2024

superlopuh Sep 29, 2024

superlopuh Oct 8, 2024

superlopuh Oct 14, 2024

superlopuh Oct 14, 2024

superlopuh Oct 14, 2024

superlopuh Oct 14, 2024

superlopuh Oct 14, 2024

superlopuh Oct 14, 2024

superlopuh Oct 14, 2024

superlopuh Oct 14, 2024

		arg_attrs = getattr(func_op, "arg_attrs")
		args = getattr(func_op, "args")

	func.func public @main(%arg0: tensor<2x3xf32>, %arg1: tensor<3x4xf32>, %arg2: tensor<2x4xf32> {tf.aliasing_output = 0 : i32}) -> (tensor<2x4xf32>) {
	func.func public @one_arg(%arg0: tensor<2x3xf32>, %arg1: tensor<3x4xf32>, %arg2: tensor<2x4xf32> {tf.aliasing_output = 0 : i32}) -> (tensor<2x4xf32>) {

	TensorMemrefInferenceConstraint("T", AnyOf([TensorType, UnrankedTensorType]))
	TensorMemrefInferenceConstraint("T", AnyTensorTypeConstr \| AnyUnrankedTensorTypeConstr)

	for inp, attr in zip(func_op.args, func_op.arg_attrs)
	for inp, attr in zip(func_op.args, func_op.arg_attrs, strict=True)

	if func_op is None or type(func_op) is not FuncOp:
	assert isinstance(func_op, FuncOp)

[WIP]: Jax tensor donation #3224

Are you sure you want to change the base?

[WIP]: Jax tensor donation #3224

Conversation

mamanain commented Sep 27, 2024

codecov bot commented Sep 27, 2024 • edited Loading

Codecov Report

superlopuh 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

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

codecov bot commented Sep 27, 2024 •

edited

Loading