Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix missing callbacks #57

Closed
wants to merge 19 commits into from
Closed

Fix missing callbacks #57

wants to merge 19 commits into from

Conversation

tatianacv
Copy link
Member

@tatianacv tatianacv commented Nov 29, 2023

Fixes wala#92.

@tatianacv
Copy link
Member Author

For the first test tf2_test_callbacks.py, the CG is:

Node: synthetic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeRootMethod()V > Context: Everywhere
 - invokestatic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeWorldClinit()V >@0
     -> Node: synthetic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeWorldClinit()V > Context: Everywhere
 - invokevirtual < PythonLoader, Lscript tf2_test_callbacks.py, do()LRoot; >@2
     -> Node: <Code body of function Lscript tf2_test_callbacks.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ]

Node: synthetic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeWorldClinit()V > Context: Everywhere

Node: <Code body of function Lscript tf2_test_callbacks.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ]
 - invokestatic < PythonLoader, Ltensorflow, import()Ltensorflow; >@90
     -> Node: synthetic < PythonLoader, Ltensorflow, import()Ltensorflow; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@90 ]
 - JSCall@104
     -> Node: synthetic < PythonLoader, Ltensorflow/distribute/MirroredStrategy, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@104 ]
 - JSCall@111
     -> Node: synthetic < PythonLoader, Ltensorflow/distribute/distribute_datasets_from_function/distribute_datasets_from_function, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@111 ]

Node: synthetic < PythonLoader, Ltensorflow, import()Ltensorflow; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@90 ]

Node: synthetic < PythonLoader, Ltensorflow/distribute/MirroredStrategy, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@104 ]

Node: synthetic < PythonLoader, Ltensorflow/distribute/distribute_datasets_from_function/distribute_datasets_from_function, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@111 ]

For the second test tf2_test_callbacks2.py, the CG is:

Node: synthetic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeRootMethod()V > Context: Everywhere
 - invokestatic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeWorldClinit()V >@0
     -> Node: synthetic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeWorldClinit()V > Context: Everywhere
 - invokevirtual < PythonLoader, Lscript tf2_test_callbacks2.py, do()LRoot; >@2
     -> Node: <Code body of function Lscript tf2_test_callbacks2.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ]

Node: synthetic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeWorldClinit()V > Context: Everywhere

Node: <Code body of function Lscript tf2_test_callbacks2.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ]
 - invokestatic < PythonLoader, Ltensorflow, import()Ltensorflow; >@90
     -> Node: synthetic < PythonLoader, Ltensorflow, import()Ltensorflow; > Context: CallStringContext: [ script tf2_test_callbacks2.py.do()LRoot;@90 ]
 - JSCall@96
     -> Node: synthetic < PythonLoader, Ltensorflow/distribute/MirroredStrategy, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks2.py.do()LRoot;@96 ]
 - JSCall@99
     -> Node: synthetic < PythonLoader, Ltensorflow/functions/constant, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks2.py.do()LRoot;@99 ]
 - JSCall@108
     -> Node: synthetic < PythonLoader, Ltensorflow/distribute/run/run, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks2.py.do()LRoot;@108 ]

Node: synthetic < PythonLoader, Ltensorflow, import()Ltensorflow; > Context: CallStringContext: [ script tf2_test_callbacks2.py.do()LRoot;@90 ]

Node: synthetic < PythonLoader, Ltensorflow/functions/constant, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks2.py.do()LRoot;@99 ]
 - invokevirtual < PythonLoader, LRoot, read_data()LRoot; >@0
     -> Node: synthetic < PythonLoader, Ltensorflow/functions/constant, read_data()LRoot; > Context: CallStringContext: [ tensorflow.functions.constant.do()LRoot;@0 ]

Node: synthetic < PythonLoader, Ltensorflow/functions/constant, read_data()LRoot; > Context: CallStringContext: [ tensorflow.functions.constant.do()LRoot;@0 ]

Node: synthetic < PythonLoader, Ltensorflow/distribute/MirroredStrategy, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks2.py.do()LRoot;@96 ]

Node: synthetic < PythonLoader, Ltensorflow/distribute/run/run, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks2.py.do()LRoot;@108 ]

@tatianacv
Copy link
Member Author

As we can see from the CGs, the functions being called are not showing up in the CG. Only the APIs that are being used.

@tatianacv
Copy link
Member Author

I ran an example on the TF Estimator referenced in the issue (wala#92 (comment)), and the CG is similar to what we see here.

@khatchad khatchad changed the title Missing Callbacks Fix missing callbacks Nov 29, 2023
@tatianacv
Copy link
Member Author

The CG for the first test:

Node: synthetic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeRootMethod()V > Context: Everywhere
 - invokestatic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeWorldClinit()V >@0
     -> Node: synthetic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeWorldClinit()V > Context: Everywhere
 - invokevirtual < PythonLoader, Lscript tf2_test_callbacks.py, do()LRoot; >@2
     -> Node: <Code body of function Lscript tf2_test_callbacks.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ]

Node: synthetic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeWorldClinit()V > Context: Everywhere

Node: <Code body of function Lscript tf2_test_callbacks.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ]
 - invokestatic < PythonLoader, Ltensorflow, import()Ltensorflow; >@90
     -> Node: synthetic < PythonLoader, Ltensorflow, import()Ltensorflow; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@90 ]
 - JSCall@104
     -> Node: synthetic < PythonLoader, Ltensorflow/distribute/MirroredStrategy, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@104 ]
 - JSCall@111
     -> Node: synthetic < PythonLoader, Ltensorflow/distribute/distribute_datasets_from_function/distribute_datasets_from_function, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@111 ]

Node: synthetic < PythonLoader, Ltensorflow, import()Ltensorflow; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@90 ]

Node: synthetic < PythonLoader, Ltensorflow/distribute/MirroredStrategy, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@104 ]

Node: synthetic < PythonLoader, Ltensorflow/distribute/distribute_datasets_from_function/distribute_datasets_from_function, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@111 ]
 - invokevirtual < PythonLoader, LRoot, do()LRoot; >@3
     -> Node: <Code body of function Lscript tf2_test_callbacks.py/get_dataset> Context: CallStringContext: [ tensorflow.distribute.distribute_datasets_from_function.distribute_datasets_from_function.do()LRoot;@3 ]

Node: <Code body of function Lscript tf2_test_callbacks.py/get_dataset> Context: CallStringContext: [ tensorflow.distribute.distribute_datasets_from_function.distribute_datasets_from_function.do()LRoot;@3 ]

@tatianacv
Copy link
Member Author

The CG for the second test:

[FINE] Call graph:
Node: synthetic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeRootMethod()V > Context: Everywhere
 - invokestatic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeWorldClinit()V >@0
     -> Node: synthetic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeWorldClinit()V > Context: Everywhere
 - invokevirtual < PythonLoader, Lscript tf2_test_callbacks2.py, do()LRoot; >@2
     -> Node: <Code body of function Lscript tf2_test_callbacks2.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ]

Node: synthetic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeWorldClinit()V > Context: Everywhere

Node: <Code body of function Lscript tf2_test_callbacks2.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ]
 - invokestatic < PythonLoader, Ltensorflow, import()Ltensorflow; >@90
     -> Node: synthetic < PythonLoader, Ltensorflow, import()Ltensorflow; > Context: CallStringContext: [ script tf2_test_callbacks2.py.do()LRoot;@90 ]
 - JSCall@96
     -> Node: synthetic < PythonLoader, Ltensorflow/distribute/MirroredStrategy, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks2.py.do()LRoot;@96 ]
 - JSCall@99
     -> Node: synthetic < PythonLoader, Ltensorflow/functions/constant, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks2.py.do()LRoot;@99 ]
 - JSCall@108
     -> Node: synthetic < PythonLoader, Ltensorflow/distribute/run/run, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks2.py.do()LRoot;@108 ]

Node: synthetic < PythonLoader, Ltensorflow, import()Ltensorflow; > Context: CallStringContext: [ script tf2_test_callbacks2.py.do()LRoot;@90 ]

Node: synthetic < PythonLoader, Ltensorflow/functions/constant, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks2.py.do()LRoot;@99 ]
 - invokevirtual < PythonLoader, LRoot, read_data()LRoot; >@0
     -> Node: synthetic < PythonLoader, Ltensorflow/functions/constant, read_data()LRoot; > Context: CallStringContext: [ tensorflow.functions.constant.do()LRoot;@0 ]

Node: synthetic < PythonLoader, Ltensorflow/functions/constant, read_data()LRoot; > Context: CallStringContext: [ tensorflow.functions.constant.do()LRoot;@0 ]

Node: synthetic < PythonLoader, Ltensorflow/distribute/MirroredStrategy, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks2.py.do()LRoot;@96 ]

Node: synthetic < PythonLoader, Ltensorflow/distribute/run/run, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks2.py.do()LRoot;@108 ]
 - invokevirtual < PythonLoader, LRoot, do()LRoot; >@3
     -> Node: <Code body of function Lscript tf2_test_callbacks2.py/replica_fn> Context: CallStringContext: [ tensorflow.distribute.run.run.do()LRoot;@3 ]

Node: <Code body of function Lscript tf2_test_callbacks2.py/replica_fn> Context: CallStringContext: [ tensorflow.distribute.run.run.do()LRoot;@3 ]

@tatianacv tatianacv marked this pull request as ready for review November 30, 2023 02:57
@khatchad
Copy link
Member

BTW, these are examples of ones that are missing. Are there more that we need that are obvious?

@tatianacv
Copy link
Member Author

BTW, these are examples of ones that are missing. Are there more that we need that are obvious?

Actually it is MirrorStrategy or TPUStrategy. This is where they are initializing strategy with method create_distribute_strategy
(https://github.com/mead-ml/mead-baseline/blob/5d7632bb151c2d09501ebf49f36ba8c4204df4c8/mead/api_examples/pretrain_discrim_tf.py#L278) and that method create_distribute_strategy uses MirrorStrategy or TPUStrategy
https://github.com/mead-ml/mead-baseline/blob/5d7632bb151c2d09501ebf49f36ba8c4204df4c8/layers/eight_mile/tf/layers.py#L121

@khatchad
Copy link
Member

The CG for the first test:

Node: synthetic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeRootMethod()V > Context: Everywhere
 - invokestatic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeWorldClinit()V >@0
     -> Node: synthetic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeWorldClinit()V > Context: Everywhere
 - invokevirtual < PythonLoader, Lscript tf2_test_callbacks.py, do()LRoot; >@2
     -> Node: <Code body of function Lscript tf2_test_callbacks.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ]

Node: synthetic < PythonLoader, Lcom/ibm/wala/FakeRootClass, fakeWorldClinit()V > Context: Everywhere

Node: <Code body of function Lscript tf2_test_callbacks.py> Context: CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ]
 - invokestatic < PythonLoader, Ltensorflow, import()Ltensorflow; >@90
     -> Node: synthetic < PythonLoader, Ltensorflow, import()Ltensorflow; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@90 ]
 - JSCall@104
     -> Node: synthetic < PythonLoader, Ltensorflow/distribute/MirroredStrategy, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@104 ]
 - JSCall@111
     -> Node: synthetic < PythonLoader, Ltensorflow/distribute/distribute_datasets_from_function/distribute_datasets_from_function, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@111 ]

Node: synthetic < PythonLoader, Ltensorflow, import()Ltensorflow; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@90 ]

Node: synthetic < PythonLoader, Ltensorflow/distribute/MirroredStrategy, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@104 ]

Node: synthetic < PythonLoader, Ltensorflow/distribute/distribute_datasets_from_function/distribute_datasets_from_function, do()LRoot; > Context: CallStringContext: [ script tf2_test_callbacks.py.do()LRoot;@111 ]
 - invokevirtual < PythonLoader, LRoot, do()LRoot; >@3
     -> Node: <Code body of function Lscript tf2_test_callbacks.py/get_dataset> Context: CallStringContext: [ tensorflow.distribute.distribute_datasets_from_function.distribute_datasets_from_function.do()LRoot;@3 ]

Node: <Code body of function Lscript tf2_test_callbacks.py/get_dataset> Context: CallStringContext: [ tensorflow.distribute.distribute_datasets_from_function.distribute_datasets_from_function.do()LRoot;@3 ]

Could we have the diffs of these?

@khatchad
Copy link
Member

Here's the diff of your CGs after some editing:

29a30,33
>  - invokevirtual < PythonLoader, LRoot, do()LRoot; >@3
>      -> Node: <Code body of function Lscript tf2_test_callbacks2.py/replica_fn> Context: CallStringContext: [ tensorflow.distribute.run.run.do()LRoot;@3 ]
> 
> Node: <Code body of function Lscript tf2_test_callbacks2.py/replica_fn> Context: CallStringContext: [ tensorflow.distribute.run.run.do()LRoot;@3 ]

This is certainly promising, but the tensor analysis still isn't picking up the tensor parameters. My guess is that the argument (in the second test; I'm not looking at the first) isn't being connected to the parameter in the XML.

@khatchad
Copy link
Member

It would also be helpful to see the IR and pointer analyses.

@khatchad
Copy link
Member

It may also be helpful to have a function that isn't a call back in the example and that directly takes a tensor parameter.

I will also say that if you simply return tf.constant() in the first test from the callback, you may seem some results. I would remove the parameter, especially since it looks like it's unused.

Copy link
Member

@khatchad khatchad left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tensor parameters don't look right.

@tatianacv tatianacv marked this pull request as draft November 30, 2023 20:42

<package name="tensorflow/distribute/run">
<class name="run" allocatable="true">
<method name="do" descriptor="()LRoot;" numArgs="2" paramNames="self fn">
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looking at https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy#run, run() takes also args=(). Looks like that needs to be passed to the callback.

@khatchad
Copy link
Member

khatchad commented Dec 1, 2023

This is really an interesting case with the Strategy.run(). The given callback can take a tuple of tensors, but Strategy.run() will send a single tensor to the callback.

@khatchad
Copy link
Member

khatchad commented Dec 1, 2023

This is really an interesting case with the Strategy.run(). The given callback can take a tuple of tensors, but Strategy.run() will send a single tensor to the callback.

(Quoting my own message here) I don't know if that's possible to say in the XML summary. We may need some Java code to unpack that.

<new def="x" class="Lobject" />
<putfield class="LRoot" field="MirroredStrategy" fieldType="LRoot" ref="self" value="x" />
<putfield class="LRoot" field="$callback" fieldType="LRoot" ref="x" value="fn" />
<call class="LRoot" name="do" descriptor="()LRoot;" type="virtual" arg0="fn" arg1="1" numArgs="2" def="xx" />
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why does arg1 take 1 here?

@@ -786,6 +792,58 @@
</class>
</package>

<package name="tensorflow/distribute">
<class name="MirroredStrategy" allocatable="true">
<method name="do" descriptor="()LRoot;" numArgs="3" paramNames="self function">
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why are there three arguments but only two parameter names?

@khatchad
Copy link
Member

khatchad commented Dec 1, 2023

OK. I think I made some kind of progress here. I'm seeing:

[Node: <Code body of function Lscript tf2_test_callbacks2.py/replica_fn> Context: CallStringContext: [ tensorflow.distribute.run.run.do()LRoot;@4 ], v2] --> [SITE_IN_NODE{<Code body of function Lscript tf2_test_callbacks2.py>:Ltuple in CallStringContext: [ com.ibm.wala.FakeRootClass.fakeRootMethod()V@2 ]}]

So, it's a tuple of tensors being sent to replica_fn, which is progress but not quite what we want.

@tatianacv Can you verify we don't have this in the PA without my latest changes?

@khatchad khatchad self-assigned this Dec 3, 2023
@khatchad khatchad added the bug Something isn't working label Dec 3, 2023
@khatchad
Copy link
Member

khatchad commented Dec 4, 2023

Superseded by #61.

@khatchad khatchad closed this Dec 4, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Missing callbacks
2 participants