Insert TBAA pointer tag for AlignedPointer and dealloc_helper Indices #1196

tzunghanjuang · 2024-10-09T18:47:39Z

Context:

Update the if-condition in TBAAPatterns so that it can insert pointer tag to Index. (Index can be treated as Pointer.)
We check if the index comes from ExtractAlignedPointerAsIndexOp or is located in dealloc_helper created by the deallocation pass. However, even with this fix, we still get double free error related Enzyme (#1167 (comment)).

Trace:
The following example has a integer tag (!tbaa !4) for Index. It should be changed to the pointer tag (!tbaa !6).

define internal void @loss.cloned(ptr %0, ptr %1, i64 %2, i64 %3, i64 %4, ptr nocapture readnone %5, ptr nocapture writeonly %6, i64 %7) {
  %9 = tail call ptr @_mlir_memref_to_llvm_alloc(i64 0)
  %10 = call ptr @_mlir_memref_to_llvm_alloc(i64 24)
  %11 = ptrtoint ptr %9 to i64
  store i64 %11, ptr %10, align 4, !tbaa !4
  ret void
}

Otherwise, this error will be thrown.

error: <unknown>:0:0: in function loss.cloned void (ptr, ptr, i64, i64, i64, ptr, ptr, i64): Enzyme: <analysis>
ptr %0: {[-1]:Pointer}, intvals: {}
ptr %1: {[-1]:Pointer}, intvals: {}
i64 %2: {[-1]:Integer}, intvals: {}
i64 %3: {[-1]:Integer}, intvals: {}
i64 %4: {[-1]:Integer}, intvals: {}
ptr %5: {[-1]:Pointer}, intvals: {}
ptr %6: {[-1]:Pointer}, intvals: {}
i64 %7: {[-1]:Integer}, intvals: {}
  %10 = call ptr @_mlir_memref_to_llvm_alloc(i64 24): {[-1]:Pointer, [-1,0]:Integer, [-1,1]:Integer, [-1,2]:Integer, [-1,3]:Integer, [-1,4]:Integer, [-1,5]:Integer, [-1,6]:Integer, [-1,7]:Integer}, intvals: {}
  %11 = ptrtoint ptr %9 to i64: {[-1]:Integer}, intvals: {}
  %9 = tail call ptr @_mlir_memref_to_llvm_alloc(i64 0): {[-1]:Pointer}, intvals: {}
</analysis>
Illegal updateAnalysis prev:{[-1]:Integer} new: {[-1]:Pointer}
val:   %11 = ptrtoint ptr %9 to i64 origin=  %11 = ptrtoint ptr %9 to i64

github-actions · 2024-10-09T19:03:37Z

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md on your branch with:

A one-to-two sentence description of the change. You may include a small working example for new features.
A link back to this PR.
Your name (or GitHub username) in the contributors section.

erick-xanadu · 2024-10-09T19:06:38Z

@tzunghanjuang please also test this in your branch which fails and report back here :) Thanks!

mlir/lib/Catalyst/Transforms/TBAAPatterns.cpp

tzunghanjuang · 2024-10-09T19:30:26Z

Failing tests in #1167: test_tensor_measure, test_multi_measure, test_qnode_different_returns, and test_multi_arg_multi_result from test_gradient_postprocessing.py.

All the error messages are sth like this.

freeing without malloc   %5 = tail call ptr @_mlir_memref_to_llvm_alloc(i64 %.idx) #7
freeing without malloc   %5 = tail call ptr @_mlir_memref_to_llvm_alloc(i64 %.idx) #7
freeing without malloc   %147 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %6, 0
freeing without malloc   %7 = tail call ptr @_mlir_memref_to_llvm_alloc(i64 72) #7
freeing without malloc   %12 = tail call ptr @_mlir_memref_to_llvm_alloc(i64 72) #7
freeing without malloc   %30 = tail call ptr @_mlir_memref_to_llvm_alloc(i64 72) #7
freeing without malloc   %35 = tail call ptr @_mlir_memref_to_llvm_alloc(i64 72) #7
free(): double free detected in tcache 2
[1]    784711 IOT instruction (core dumped)  python a_test/test_tensor_measure.py

erick-xanadu · 2024-10-09T19:45:04Z

Failing tests in #1167: test_tensor_measure, test_multi_measure, test_qnode_different_returns, and test_multi_arg_multi_result from test_gradient_postprocessing.py.

@tzunghanjuang , you are saying that even with this PR, these tests in your other PR are still failing?

tzunghanjuang · 2024-10-09T19:58:39Z

@erick-xanadu Yes. These still fails with this PR.

dime10 · 2024-10-09T21:09:24Z

mlir/lib/Catalyst/Transforms/TBAAPatterns.cpp

+        // Index can be used as a pointer.
+        if (isa<IndexType>(baseType) && (isa<LLVM::StoreOp>(newOp) || isa<LLVM::LoadOp>(newOp))) {


I don't know if we should always consider index as a pointer. Rather, I would walk backwards and if the index comes from an op that we know represents a pointer (like memref.extract_aligned_pointer), we change the TBAA to pointer.
Always considering index to be a pointer sounds like it could produce a lot of errors.

I am not sure if there is any way to memerize if a memref consider from memref.extract_aligned_pointer. Maybe we need to modify the mlir side to add this information?

Edit: I added a helper method that recursively tracks prevNode until it encounters memref.extract_aligned_pointer.

erick-xanadu · 2024-10-09T21:28:39Z

I think a prerequisite for merging this PR be that it solves the issue in your other PR. @tzunghanjuang

codecov · 2024-10-09T21:39:22Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.93%. Comparing base (c0803e5) to head (ed5d99c).
Report is 6 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1196      +/-   ##
==========================================
- Coverage   97.93%   97.93%   -0.01%     
==========================================
  Files          76       77       +1     
  Lines       10893    10921      +28     
  Branches      970      971       +1     
==========================================
+ Hits        10668    10695      +27     
  Misses        179      179              
- Partials       46       47       +1

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

tzunghanjuang · 2024-10-10T00:59:57Z

The double free issue might be related to this EnzymeAD/Enzyme#2029.
Edit: Our Enzyme version already has the fix but we still have the error.

dime10 · 2024-10-11T17:18:23Z

mlir/lib/Catalyst/Transforms/TBAAPatterns.cpp

+            (isFromExtractAlignedPointerAsIndexOp(currentOp) || isInsideDeallocHelper(currentOp))) {
+            tag = tree->getTag("any pointer");
+        }
+        else {
+            isInsideDeallocHelper(currentOp);


What's the idea behind the isInsideDeallocHelper function? Is isFromExtractAlignedPointerAsIndexOp not sufficient?

Also the second call to isInsideDeallocHelper doesn't seem to do anything.

The deallocation pass inserts delloc_helper function and it does not come with any ExtractAlignedPointerAsIndexOp.

The index pointers (like %66 at the below) are fed to dealloc_helper (%21). So the store and load inside dealloc_helper should be added with the pointer tag.

%66 = call ptr @_mlir_memref_to_llvm_alloc(i64 ptrtoint (ptr getelementptr (i1, ptr null, i32 1) to i64)) call void @dealloc_helper(ptr %54, ptr %54, i64 0, i64 3, i64 1, ptr %56, ptr %56, i64 0, i64 1, i64 1, ptr %55, ptr %55, i64 0, i64 3, i64 1, ptr %65, ptr %65, i64 0, i64 3, i64 1, ptr %66, ptr %66, i64 0, i64 1, i64 1) ... define void @dealloc_helper(ptr %0, ptr %1, i64 %2, i64 %3, i64 %4, ptr %5, ptr %6, i64 %7, i64 %8, i64 %9, ptr %10, ptr %11, i64 %12, i64 %13, i64 %14, ptr %15, ptr %16, i64 %17, i64 %18, i64 %19, ptr %20, ptr %21, i64 %22, i64 %23, i64 %24) { ... 45: ; preds = %42 %46 = getelementptr inbounds i1, ptr %21, i64 %43 store i1 false, ptr %46, align 1, !tbaa !6 %47 = add i64 %43, 1 br label %42 ... }

I think this dealloc_helper handles index pointers, so we should insert pointer tags inside it.

I see, so actually we would need to follow all uses of the buffers (memrefs) where the index-typed pointers are stored into. Tricky indeed.

rmoyard · 2024-10-11T17:41:18Z

mlir/lib/Catalyst/Transforms/TBAAPatterns.cpp

-        tag = tree->getTag("int");
+        // Index can be used as a pointer.
+        if (isa<IndexType>(baseType) &&
+            (isFromExtractAlignedPointerAsIndexOp(currentOp) || isInsideDeallocHelper(currentOp))) {


Can you extract the logic isFromExtractAlignedPointerAsIndexOp(currentOp) || isInsideDeallocHelper(currentOp) from set tags? You should be able to determine the type before the setting it.

Yes, I have moved this logic out.

rmoyard · 2024-10-11T17:42:36Z

mlir/lib/Catalyst/Transforms/TBAAPatterns.cpp

@@ -72,7 +108,7 @@ struct MemrefLoadTBAARewritePattern : public ConvertOpToLLVMPattern<memref::Load
            loadOp.getNontemporal());

        if (isAnyOf<IndexType, IntegerType, FloatType, MemRefType>(baseType)) {


You can add the a function that check that is index and isFromExtractAlignedPointerAsIndexOp(currentOp) || isInsideDeallocHelper(currentOp) then you return ptr or index and do not need to change the settag

Tzung-Han Juang added 4 commits October 9, 2024 13:50

Make TBAA pass not set tag to Index with Store and Load

3a271b1

Update TBBA test

23d9fb7

TBAA pass insert pointer tag to index (with Load/Store)

53c4ed7

Reformatting

63eeb6f

tzunghanjuang requested review from erick-xanadu and rmoyard October 9, 2024 19:03

rmoyard reviewed Oct 9, 2024

View reviewed changes

mlir/lib/Catalyst/Transforms/TBAAPatterns.cpp Outdated Show resolved Hide resolved

mlir/lib/Catalyst/Transforms/TBAAPatterns.cpp Outdated Show resolved Hide resolved

dime10 reviewed Oct 9, 2024

View reviewed changes

Tzung-Han Juang added 2 commits October 9, 2024 17:23

Revert changes on TBAA test

bba5375

Apply pointer tag if prevNode is PtrToIntOp

100f736

Tzung-Han Juang added 2 commits October 10, 2024 15:41

Check if index comes from aligned pointer or delloc helper

0cfe2fa

Improve helpers' name

982ba01

tzunghanjuang changed the title ~~TBAA insserts pointer tag to Index (with Load or Store)~~ Insert TBAA pointer tag for AlignedPointer and dealloc_helper Indices Oct 10, 2024

dime10 reviewed Oct 11, 2024

View reviewed changes

Cleanup

a3bea1a

rmoyard reviewed Oct 11, 2024

View reviewed changes

Move index pointer case out of setTag

ed5d99c

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Insert TBAA pointer tag for AlignedPointer and dealloc_helper Indices #1196

Insert TBAA pointer tag for AlignedPointer and dealloc_helper Indices #1196

tzunghanjuang commented Oct 9, 2024 •

edited

Loading

github-actions bot commented Oct 9, 2024

erick-xanadu commented Oct 9, 2024

tzunghanjuang commented Oct 9, 2024 •

edited

Loading

erick-xanadu commented Oct 9, 2024

tzunghanjuang commented Oct 9, 2024

dime10 Oct 9, 2024 •

edited

Loading

tzunghanjuang Oct 9, 2024 •

edited

Loading

erick-xanadu commented Oct 9, 2024

codecov bot commented Oct 9, 2024 •

edited

Loading

tzunghanjuang commented Oct 10, 2024 •

edited

Loading

dime10 Oct 11, 2024

tzunghanjuang Oct 11, 2024

tzunghanjuang Oct 11, 2024

dime10 Oct 11, 2024 •

edited

Loading

rmoyard Oct 11, 2024

tzunghanjuang Oct 11, 2024

rmoyard Oct 11, 2024

		// Index can be used as a pointer.
		if (isa<IndexType>(baseType) && (isa<LLVM::StoreOp>(newOp) \|\| isa<LLVM::LoadOp>(newOp))) {

		@@ -72,7 +108,7 @@ struct MemrefLoadTBAARewritePattern : public ConvertOpToLLVMPattern<memref::Load
		loadOp.getNontemporal());

		if (isAnyOf<IndexType, IntegerType, FloatType, MemRefType>(baseType)) {

Insert TBAA pointer tag for AlignedPointer and dealloc_helper Indices #1196

Are you sure you want to change the base?

Insert TBAA pointer tag for AlignedPointer and dealloc_helper Indices #1196

Conversation

tzunghanjuang commented Oct 9, 2024 • edited Loading

github-actions bot commented Oct 9, 2024

erick-xanadu commented Oct 9, 2024

tzunghanjuang commented Oct 9, 2024 • edited Loading

erick-xanadu commented Oct 9, 2024

tzunghanjuang commented Oct 9, 2024

dime10 Oct 9, 2024 • edited Loading

Choose a reason for hiding this comment

tzunghanjuang Oct 9, 2024 • edited Loading

Choose a reason for hiding this comment

erick-xanadu commented Oct 9, 2024

codecov bot commented Oct 9, 2024 • edited Loading

Codecov Report

tzunghanjuang commented Oct 10, 2024 • edited Loading

dime10 Oct 11, 2024

Choose a reason for hiding this comment

tzunghanjuang Oct 11, 2024

Choose a reason for hiding this comment

tzunghanjuang Oct 11, 2024

Choose a reason for hiding this comment

dime10 Oct 11, 2024 • edited Loading

Choose a reason for hiding this comment

rmoyard Oct 11, 2024

Choose a reason for hiding this comment

tzunghanjuang Oct 11, 2024

Choose a reason for hiding this comment

rmoyard Oct 11, 2024

Choose a reason for hiding this comment

tzunghanjuang commented Oct 9, 2024 •

edited

Loading

tzunghanjuang commented Oct 9, 2024 •

edited

Loading

dime10 Oct 9, 2024 •

edited

Loading

tzunghanjuang Oct 9, 2024 •

edited

Loading

codecov bot commented Oct 9, 2024 •

edited

Loading

tzunghanjuang commented Oct 10, 2024 •

edited

Loading

dime10 Oct 11, 2024 •

edited

Loading