[SROA][TBAA] Handle shift of regular TBAA nodes

[wsmoses]

Summary

SROA shifts TBAA nodes in a way that may present a problem for !tbaa but not !tbaa.struct nodes.

[aheejin]

I don't know enough about this code to review this, but I can confirm that this patch passes my reduced example code and also the original code. Thanks!

[aheejin]

@jdoerfert Can you please take a look? One of our clients is experiencing this problem, and I don't have enough knowledge to review this area of code.. Thanks!

[aheejin]

It looks the other reviewer is not looking at this, so I'll approve this to unblock our client. I don't have enough context in this code, so I'd appreciate if other reviewers comment retroactively.
@wsmoses Thanks for the quick fix!

[jeroen.dobbelaere]

Please do not commit this yet. A ping would have been appreciated.

[aheejin]

@jeroen.dobbelaere Sorry not sure if I understand. Are you planning to review the code?

[jeroen.dobbelaere]

In D99851#2678971, @aheejin wrote:

@jeroen.dobbelaere Sorry not sure if I understand. Are you planning to review the code?

Yes, I have been doing this at the moment. Sorry for taking that long.

It indeed seems that D95826 has been shifting this part of TBAA in the wrong direction. The change in this patch also does not trigger any test failures, so this is also not tested well.

I would like to see a decent fix, meaning that the offset is shifted in the other direction. The test example should also be extended to show the case where the shift should result in different tbaa types for the resulting stores.

Probably care must be taken when this is not about a member of a struct type: https://llvm.org/docs/LangRef.html#tbaa-metadata says that for basic types, the offset must always be 0.
An extra test where for example a i64 store (or load), representing a basic 'long' type or so, is split into two i32 stores (loads) would also be appreciated.

[wsmoses]

I'm not quite sure I understand what you mean @jeroen.dobbelaere by an i64 store of a long being split to 2 i32 int stores. I presume you mean that there was originally a struct {int, int} that is currently using an i64 store to do both at the same time. I don't believe that a "long" store would have this apply since its tbaa would simply be one long, right? Moreover, this shifting is only in use in SROA at the moment so I'm not sure it would split a store of a non-aggregate type (such as an i64).

I shall update the test case. Without the patch, in my testing verifier fails, but the actual TBAA nodes should be actually defined in the check's.

My initial concern with the mechanism for adding the offset is that naiively adding the offset here may be legal as it appears that TBAA nodes must have an offset corresponding to an offset defined in the base node. This may make it possible to perform a legal shift that does not have a defined offset within the base type (and in practice on the test suite, this appeared to occur).

I'd be happy to take a further look to try to examine this but won't be able to start diving deeply in for at least another week (several imminent conference paper deadlines). If someone else would like to take over this in the interim, feel free. It sounded from @aheejin that this was somewhat time sensitive.

[aheejin]

Yes, this is time sensitive, and what I need is not optimization but just correct code generation now. @jeroen.dobbelaere, as long as this patch generates correct code, can we land this for now and let @wsmoses take another stab in a follow-up patch? This is blocking https://bugs.llvm.org/show_bug.cgi?id=49825, which is a reduced test case of our client's code.

[jeroen.dobbelaere]

llvm/test/Transforms/SROA/tbaa-subload.ll

36

@aheejin The issue I see is, that if !8 would be something like :

!8 = !{!"_ZTSZN2ax2baEMS_FvvE2an2arE3$_0", !9, i64 0, !10, i64 8, !11, i64 16}
!10 = !{!"long", !3, i64 0 }
!11 = !{!"int", !3, i64 0 }

The resulting tbaa will still be incorrect after this patch: the two stores should be referring to two different items, with different offsets.
aka, the resulting !0 = !{!1, !4, i64 8} is only valid for the first store. The second store should be something like '!99 = !{!1, !4, i64 16 }

[jeroen.dobbelaere]

FYI: I am looking into an alternative fix.

[wsmoses]

Ah ok understood. And yes I completely agree that the patch as is would fail on this case (the patch as is essentially reverts the tbaa shift to the behavior prior to https://reviews.llvm.org/D95826 for TBAA nodes while retaining the correct behavior for tbaa.struct nodes).

I completely agree with doing a proper fix, especially if @jeroen.dobbelaere you have the cycles to do so right now. Is there anything I need to do to reassign this to you? For now I remain in a constant MLIR hackathon until the PACT deadline on Monday.

[jeroen.dobbelaere]

@wsmoses I will not have a solution today, but hopefully tomorrow. I would appreciate if you could help review then. Thanks.

[wsmoses]

Will do.

[jeroen.dobbelaere]

I have been investigating the behavior of struct path tbaa and new struct path tbaa. Based on what I could find, returning MD is the correct thing to do in struct path tbaa.

My understanding is that for new struct path tbaa, it is/was the intention to support shifting, but the Verifier.cpp is not allowing this at this moment.
So for now, returning MD should also be ok at this moment.

It is also interesting to note that the { i64, i64 } struct, represented as a char with size 16, originates from a member function pointer. clang will map unhandled types to char by default. (See CodeGenTBAA.cpp: getTypeInfoHelper(..))
See D40176 and D39956 for some more info about the new struct path tbaa. IMHO, we should pick up this work, come up with a decent specification and provide in-tree documentation.
(@kosarev As far as I could find, the current specification is in https://lists.llvm.org/pipermail/llvm-dev/2017-November/118748.html )

I'll bring this up in next weeks LLVM AA Tech Call.

[aheejin]

Thanks! @wsmoses, can you commit the patch?

llvm/lib/Analysis/TypeBasedAliasAnalysis.cpp

 MDNode *AAMDNodes::ShiftTBAA(MDNode *MD, size_t Offset) {
   // Fast path if there's no offset
   if (Offset == 0)
     return MD;
   // Fast path if there's no path tbaa node (and thus scalar)
   if (!isStructPathTBAA(MD))
     return MD;
 
-  TBAAStructTagNode Tag(MD);
-  SmallVector<Metadata *, 5> Sub;
-  Sub.push_back(MD->getOperand(0));
-  Sub.push_back(MD->getOperand(1));
-  ConstantInt *InnerOffset = mdconst::extract<ConstantInt>(MD->getOperand(2));
-
-  if (Tag.isNewFormat()) {
-    ConstantInt *InnerSize = mdconst::extract<ConstantInt>(MD->getOperand(3));
-
-    if (InnerOffset->getZExtValue() + InnerSize->getZExtValue() <= Offset) {
-      return nullptr;
-    }
-
-    uint64_t NewSize = InnerSize->getZExtValue();
-    uint64_t NewOffset = InnerOffset->getZExtValue() - Offset;
-    if (InnerOffset->getZExtValue() < Offset) {
-      NewOffset = 0;
-      NewSize -= Offset - InnerOffset->getZExtValue();
-    }
-
-    Sub.push_back(ConstantAsMetadata::get(
-        ConstantInt::get(InnerOffset->getType(), NewOffset)));
-
-    Sub.push_back(ConstantAsMetadata::get(
-        ConstantInt::get(InnerSize->getType(), NewSize)));
-
-    // immutable type
-    if (MD->getNumOperands() >= 5)
-      Sub.push_back(MD->getOperand(4));
-  } else {
-    if (InnerOffset->getZExtValue() < Offset)
-      return nullptr;
-
-    Sub.push_back(ConstantAsMetadata::get(ConstantInt::get(
-        InnerOffset->getType(), InnerOffset->getZExtValue() - Offset)));
-
-    // immutable type
-    if (MD->getNumOperands() >= 4)
-      Sub.push_back(MD->getOperand(3));
-  }
-  return MDNode::get(MD->getContext(), Sub);
+  // The correct behavior here is to add the offset into the TBAA
+  // struct node offset. The base type, however may not have defined
+  // a type at this additional offset, resulting in errors. Since
+  // this method is only used within a given load/store access
+  // the offset provided is only used to subdivide the previous load
+  // maintaining the validity of the previous TBAA.
+  //
+  // This, however, should be revisited in the future.
+  return MD;
 }
 
 MDNode *AAMDNodes::ShiftTBAAStruct(MDNode *MD, size_t Offset) {
   // Fast path if there's no offset
   if (Offset == 0)
     return MD;
   SmallVector<Metadata *, 3> Sub;
   for (size_t i = 0, size = MD->getNumOperands(); i < size; i += 3) {

llvm/test/Transforms/SROA/tbaa-subload.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -sroa %s | FileCheck %s
+
+; This should not crash
+
+%class.ar = type { i8 }
+%class.anon = type <{ %class.ar, [7 x i8], { i64, i64 } }>
+
+define void @caller() {
+; CHECK-LABEL: @caller(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[AGG:%.*]] = alloca [[CLASS_ANON:%.*]], align 8
+; CHECK-NEXT:    [[OFF:%.*]] = getelementptr inbounds [[CLASS_ANON]], %class.anon* [[AGG]], i32 0, i32 2
+; CHECK-NEXT:    [[DOTFCA_0_GEP:%.*]] = getelementptr inbounds { i64, i64 }, { i64, i64 }* [[OFF]], i32 0, i32 0
+; CHECK-NEXT:    store i64 1, i64* [[DOTFCA_0_GEP]], align 8, !tbaa [[TBAA0:![0-9]+]]
+; CHECK-NEXT:    [[DOTFCA_1_GEP:%.*]] = getelementptr inbounds { i64, i64 }, { i64, i64 }* [[OFF]], i32 0, i32 1
+; CHECK-NEXT:    store i64 2, i64* [[DOTFCA_1_GEP]], align 8, !tbaa [[TBAA0]]
+; CHECK-NEXT:    call void @use(%class.anon* [[AGG]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %agg = alloca %class.anon, align 8
+  %off = getelementptr inbounds %class.anon, %class.anon* %agg, i32 0, i32 2
+  store { i64, i64 } { i64 1, i64 2 }, { i64, i64 }* %off, align 8, !tbaa !7
+  call void @use(%class.anon* %agg)
+  ret void
+}
+
+declare void @use(%class.anon* %this)
+
+!3 = !{!"omnipotent char", !4, i64 0}
+!4 = !{!"Simple C++ TBAA"}
+
+!7 = !{!8, !3, i64 8}
+!8 = !{!"_ZTSZN2ax2baEMS_FvvE2an2arE3$_0", !9, i64 0, !3, i64 8}
+!9 = !{!"_ZTS2ar"}