InferAddressSpaces: Handle icmp

[jlebar]

Well this one is interesting.

I think I see what you're getting at here; if we can ASC away all uses of a generic pointer except the icmp, then keeping around the generic pointers in icmp is just unnecessary register pressure.

But at least on NVPTX, I think this is only safe if we can transform *both* pointers in the icmp, and moreover only if both pointers are transformed to the same AS.

[arsenm]

In D29213#658543, @jlebar wrote:

Well this one is interesting.

I think I see what you're getting at here; if we can ASC away all uses of a generic pointer except the icmp, then keeping around the generic pointers in icmp is just unnecessary register pressure.

But at least on NVPTX, I think this is only safe if we can transform *both* pointers in the icmp, and moreover only if both pointers are transformed to the same AS.

The IR type system requires both be the same type, so this is done here. Also I just noticed I forgot to update the TODO saying to do this

[jlebar]

lib/Transforms/Scalar/InferAddressSpaces.cpp

747

I don't see where we check that NewV and OtherNewV have the same AS. Shouldn't it be right here?

Sorry if it's obvious; I'm breaking my rule against reviewing patches late at night, and may be embarrassed in the morning.

[arsenm]

lib/Transforms/Scalar/InferAddressSpaces.cpp

747

I was confused about what you were saying, but the check turns out to have been added in the patch I forgot to squash which fixed that bug

[jlebar]

Just so I know how to read these tests: AS(4) is the flat AS?

lib/Transforms/Scalar/InferAddressSpaces.cpp

759

Can we have a test for icmp x, undef and icmp undef, x?

762

If KOtherSrc->AS() == NewAS, then isn't this ASC redundant? It seems we should only need it if KOtherSrc is undef, and in that case, can we just substitute a new undef value, instead of ASC'ing the old undef?

[arsenm]

In D29213#658940, @jlebar wrote:

Just so I know how to read these tests: AS(4) is the flat AS?

Yes

lib/Transforms/Scalar/InferAddressSpaces.cpp

759

I can add another test for undef, x, but there's no point in trying to handle these in the code since these are canonicalized to constants in the RHS

762

I made this too conservative by adding the new check here, so this is backwards. The addrspace check is redundant with the cast. We need to cast the constants, otherwise it's making an assumption about the bit values of the pointer. For undef I think replacing the type directly would be OK, but the ConstantExpr cast should do the right thing (and I think this will be necessary in the future with freeze anyway)

[arsenm]

Don't check type for constants, just cast them

[jlebar]

lib/Transforms/Scalar/InferAddressSpaces.cpp

759

sgtm

779

Since this is now generic, we probably shouldn't say "shared" here.

784

Is this comment now out of date? I'm confused about "instead of currently".

Would also prefer for the comment to start with prose. For example:

// If we can infer that both pointers are in the same addrspace, transform e.g.
//   %cmp = icmp eq float*%p, %q
// into
//   %cmp = icmp eq float addrspace(3)* %new_p, %new_q

Is that all we need to say here?

[arsenm]

Additional checks are necessary with constant expressions

[arsenm]

r293593

lib/Transforms/Scalar/InferAddressSpaces.cpp

if (LoadInst *LI = dyn_cast<LoadInst>(&I)) {
 
       // Handle 2nd operand for memcpy/memmove.
       if (MemTransferInst *MTI = dyn_cast<MemTransferInst>(MI)) {
         appendsGenericAddressExpressionToPostorderStack(
           MTI->getRawSource(), &PostorderStack, &Visited);
       }
     } else if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I)) {
       collectRewritableIntrinsicOperands(II, &PostorderStack, &Visited);
+    } else if (ICmpInst *Cmp = dyn_cast<ICmpInst>(&I)) {
+      // FIXME: Handle vectors of pointers
+      if (Cmp->getOperand(0)->getType()->isPointerTy()) {
+        appendsGenericAddressExpressionToPostorderStack(
+          Cmp->getOperand(0), &PostorderStack, &Visited);
+        appendsGenericAddressExpressionToPostorderStack(
+          Cmp->getOperand(1), &PostorderStack, &Visited);
+      }
     }
   }
 
   std::vector<Value *> Postorder; // The resultant postorder.
   while (!PostorderStack.empty()) {
     // If the operands of the expression on the top are already explored,
     // adds that expression to the resultant postorder.
     if (PostorderStack.back().second) {
for (I = V->use_begin(), E = V->use_end(); I != E; ) {
       while (I != E && I->getUser() == CurUser) {
         ++I;
       }
 
       // Handle more complex cases like intrinsic that need to be remangled.
       if (handleComplexPtrUse(*CurUser, V, NewV))
         continue;
 
-      if (isa<Instruction>(U.getUser())) {
+      if (isa<Instruction>(CurUser)) {
+        if (ICmpInst *Cmp = dyn_cast<ICmpInst>(CurUser)) {
+
+          unsigned NewAS = NewV->getType()->getPointerAddressSpace();
+          int SrcIdx = U.getOperandNo();
+          int OtherIdx = (SrcIdx == 0) ? 1 : 0;
+          Value *OtherSrc = Cmp->getOperand(OtherIdx);
+
+          if (Value *OtherNewV = ValueWithNewAddrSpace.lookup(OtherSrc)) {
+            if (OtherNewV->getType()->getPointerAddressSpace() == NewAS) {
+              Cmp->setOperand(OtherIdx, OtherNewV);
+              Cmp->setOperand(SrcIdx, NewV);
+              continue;
+            }
+          }
+
+          if (auto *KOtherSrc = dyn_cast<Constant>(OtherSrc)) {
+            if (KOtherSrc->getType()->getPointerAddressSpace() == NewAS ||
+                isa<UndefValue>(KOtherSrc)) {
+              Cmp->setOperand(SrcIdx, NewV);
+              Cmp->setOperand(OtherIdx,
+                ConstantExpr::getAddrSpaceCast(KOtherSrc, NewV->getType()));
+              continue;
+            }
+          }
+        }
+
         // Otherwise, replaces the use with generic(NewV).
         // TODO: Some optimization opportunities are missed. For example, in
         //   %0 = icmp eq float* %p, %q
         // if both p and q are inferred to be shared, we can rewrite %0 as
         //   %0 = icmp eq float addrspace(3)* %new_p, %new_q
         // instead of currently
         //   %generic_p = addrspacecast float addrspace(3)* %new_p to float*
         //   %generic_q = addrspacecast float addrspace(3)* %new_q to float*
         //   %0 = icmp eq float* %generic_p, %generic_q
         if (Instruction *I = dyn_cast<Instruction>(V)) {
           BasicBlock::iterator InsertPos = std::next(I->getIterator());
           while (isa<PHINode>(InsertPos))
             ++InsertPos;
           U.set(new AddrSpaceCastInst(NewV, V->getType(), "", &*InsertPos));
         } else {
           U.set(ConstantExpr::getAddrSpaceCast(cast<Constant>(NewV),
                                                V->getType()));
         }
       }
     }
 
     if (V->use_empty())
       RecursivelyDeleteTriviallyDeadInstructions(V);
   }
 
   return true;
 }
 
 FunctionPass *llvm::createInferAddressSpacesPass() {
   return new InferAddressSpaces();
 }

test/Transforms/InferAddressSpaces/AMDGPU/icmp.ll

+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -infer-address-spaces %s | FileCheck %s
+
+; CHECK-LABEL: @icmp_flat_cmp_self(
+; CHECK: %cmp = icmp eq i32 addrspace(3)* %group.ptr.0, %group.ptr.0
+define i1 @icmp_flat_cmp_self(i32 addrspace(3)* %group.ptr.0) #0 {
+  %cast0 = addrspacecast i32 addrspace(3)* %group.ptr.0 to i32 addrspace(4)*
+  %cmp = icmp eq i32 addrspace(4)* %cast0, %cast0
+  ret i1 %cmp
+}
+
+; CHECK-LABEL: @icmp_flat_flat_from_group(
+; CHECK: %cmp = icmp eq i32 addrspace(3)* %group.ptr.0, %group.ptr.1
+define i1 @icmp_flat_flat_from_group(i32 addrspace(3)* %group.ptr.0, i32 addrspace(3)* %group.ptr.1) #0 {
+  %cast0 = addrspacecast i32 addrspace(3)* %group.ptr.0 to i32 addrspace(4)*
+  %cast1 = addrspacecast i32 addrspace(3)* %group.ptr.1 to i32 addrspace(4)*
+  %cmp = icmp eq i32 addrspace(4)* %cast0, %cast1
+  ret i1 %cmp
+}
+
+; CHECK-LABEL: @icmp_mismatch_flat_from_group_private(
+; CHECK: %1 = addrspacecast i32* %private.ptr.0 to i32 addrspace(4)*
+; CHECK: %2 = addrspacecast i32 addrspace(3)* %group.ptr.1 to i32 addrspace(4)*
+; CHECK: %cmp = icmp eq i32 addrspace(4)* %1, %2
+define i1 @icmp_mismatch_flat_from_group_private(i32* %private.ptr.0, i32 addrspace(3)* %group.ptr.1) #0 {
+  %cast0 = addrspacecast i32* %private.ptr.0 to i32 addrspace(4)*
+  %cast1 = addrspacecast i32 addrspace(3)* %group.ptr.1 to i32 addrspace(4)*
+  %cmp = icmp eq i32 addrspace(4)* %cast0, %cast1
+  ret i1 %cmp
+}
+
+ ; CHECK-LABEL: @icmp_flat_group_flat(
+; CHECK: %1 = addrspacecast i32 addrspace(3)* %group.ptr.0 to i32 addrspace(4)*
+; CHECK: %cmp = icmp eq i32 addrspace(4)* %1, %flat.ptr.1
+define i1 @icmp_flat_group_flat(i32 addrspace(3)* %group.ptr.0, i32 addrspace(4)* %flat.ptr.1) #0 {
+  %cast0 = addrspacecast i32 addrspace(3)* %group.ptr.0 to i32 addrspace(4)*
+  %cmp = icmp eq i32 addrspace(4)* %cast0, %flat.ptr.1
+  ret i1 %cmp
+}
+
+; CHECK-LABEL: @icmp_flat_flat_group(
+; CHECK: %1 = addrspacecast i32 addrspace(3)* %group.ptr.1 to i32 addrspace(4)*
+; CHECK: %cmp = icmp eq i32 addrspace(4)* %flat.ptr.0, %1
+define i1 @icmp_flat_flat_group(i32 addrspace(4)* %flat.ptr.0, i32 addrspace(3)* %group.ptr.1) #0 {
+  %cast1 = addrspacecast i32 addrspace(3)* %group.ptr.1 to i32 addrspace(4)*
+  %cmp = icmp eq i32 addrspace(4)* %flat.ptr.0, %cast1
+  ret i1 %cmp
+}
+
+; Keeping as cmp addrspace(3)* is better
+; CHECK-LABEL: @icmp_flat_to_group_cmp(
+; CHECK: %cast0 = addrspacecast i32 addrspace(4)* %flat.ptr.0 to i32 addrspace(3)*
+; CHECK: %cast1 = addrspacecast i32 addrspace(4)* %flat.ptr.1 to i32 addrspace(3)*
+; CHECK: %cmp = icmp eq i32 addrspace(3)* %cast0, %cast1
+define i1 @icmp_flat_to_group_cmp(i32 addrspace(4)* %flat.ptr.0, i32 addrspace(4)* %flat.ptr.1) #0 {
+  %cast0 = addrspacecast i32 addrspace(4)* %flat.ptr.0 to i32 addrspace(3)*
+  %cast1 = addrspacecast i32 addrspace(4)* %flat.ptr.1 to i32 addrspace(3)*
+  %cmp = icmp eq i32 addrspace(3)* %cast0, %cast1
+  ret i1 %cmp
+}
+
+; FIXME: Should be able to ask target about how to constant fold the
+; constant cast if this is OK to change if 0 is a valid pointer.
+
+; CHECK-LABEL: @icmp_group_flat_cmp_null(
+; CHECK: %1 = addrspacecast i32 addrspace(3)* %group.ptr.0 to i32 addrspace(4)*
+; CHECK: %cmp = icmp eq i32 addrspace(4)* %1, null
+define i1 @icmp_group_flat_cmp_null(i32 addrspace(3)* %group.ptr.0) #0 {
+  %cast0 = addrspacecast i32 addrspace(3)* %group.ptr.0 to i32 addrspace(4)*
+  %cmp = icmp eq i32 addrspace(4)* %cast0, null
+  ret i1 %cmp
+}
+
+; CHECK-LABEL: @icmp_mismatch_flat_group_private_cmp_null(
+; CHECK: %1 = addrspacecast i32 addrspace(3)* %group.ptr.0 to i32 addrspace(4)*
+; CHECK: %cmp = icmp eq i32 addrspace(4)* %1, addrspacecast (i32* null to i32 addrspace(4)*)
+define i1 @icmp_mismatch_flat_group_private_cmp_null(i32 addrspace(3)* %group.ptr.0) #0 {
+  %cast0 = addrspacecast i32 addrspace(3)* %group.ptr.0 to i32 addrspace(4)*
+  %cmp = icmp eq i32 addrspace(4)* %cast0, addrspacecast (i32* null to i32 addrspace(4)*)
+  ret i1 %cmp
+}
+
+; CHECK-LABEL: @icmp_group_flat_cmp_undef(
+; CHECK: %cmp = icmp eq i32 addrspace(3)* %group.ptr.0, undef
+define i1 @icmp_group_flat_cmp_undef(i32 addrspace(3)* %group.ptr.0) #0 {
+  %cast0 = addrspacecast i32 addrspace(3)* %group.ptr.0 to i32 addrspace(4)*
+  %cmp = icmp eq i32 addrspace(4)* %cast0, undef
+  ret i1 %cmp
+}
+
+; TODO: Should be handled
+; CHECK-LABEL: @icmp_flat_flat_from_group_vector(
+; CHECK: %cmp = icmp eq <2 x i32 addrspace(4)*> %cast0, %cast1
+define <2 x i1> @icmp_flat_flat_from_group_vector(<2 x i32 addrspace(3)*> %group.ptr.0, <2 x i32 addrspace(3)*> %group.ptr.1) #0 {
+  %cast0 = addrspacecast <2 x i32 addrspace(3)*> %group.ptr.0 to <2 x i32 addrspace(4)*>
+  %cast1 = addrspacecast <2 x i32 addrspace(3)*> %group.ptr.1 to <2 x i32 addrspace(4)*>
+  %cmp = icmp eq <2 x i32 addrspace(4)*> %cast0, %cast1
+  ret <2 x i1> %cmp
+}
+
+attributes #0 = { nounwind }

test/Transforms/InferAddressSpaces/AMDGPU/infer-address-space.ll

bb:
   store float %v, float addrspace(4)* %addr
   ret void
 }
 
 ; FIXME: Should be able to eliminate the cast inside the loop
 ; CHECK-LABEL: @loop(
 
 ; CHECK: %p = bitcast [10 x float] addrspace(3)* @array to float addrspace(3)*
-; CHECK: %0 = addrspacecast float addrspace(3)* %p to float addrspace(4)*
-; CHECK: %end = getelementptr float, float addrspace(4)* %0, i64 10
+; CHECK: %end = getelementptr float, float addrspace(3)* %p, i64 10
 ; CHECK: br label %loop
 
 ; CHECK: loop:                                             ; preds = %loop, %entry
 ; CHECK: %i = phi float addrspace(3)* [ %p, %entry ], [ %i2, %loop ]
 ; CHECK: %v = load float, float addrspace(3)* %i
 ; CHECK: call void @use(float %v)
 ; CHECK: %i2 = getelementptr float, float addrspace(3)* %i, i64 1
-; CHECK: %1 = addrspacecast float addrspace(3)* %i2 to float addrspace(4)*
-; CHECK: %exit_cond = icmp eq float addrspace(4)* %1, %end
+; CHECK: %exit_cond = icmp eq float addrspace(3)* %i2, %end
+
 ; CHECK: br i1 %exit_cond, label %exit, label %loop
 define void @loop() #0 {
 entry:
   %p = addrspacecast [10 x float] addrspace(3)* @array to float addrspace(4)*
   %end = getelementptr float, float addrspace(4)* %p, i64 10
   br label %loop
 
 loop:                                             ; preds = %loop, %entry