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

779

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

779

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

791

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

794

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

791

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

794

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

791

sgtm

808

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

813

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

private:
 
   // Tries to infer the specific address space of each address expression in
   // Postorder.
   void inferAddressSpaces(const std::vector<Value *> &Postorder,
                           ValueToAddrSpaceMapTy *InferredAddrSpace) const;
 
   bool handleComplexPtrUse(User &U, Value *OldV, Value *NewV) const;
 
+  bool isSafeToCastConstAddrSpace(Constant *C, unsigned NewAS) const;
+
   // Changes the generic address expressions in function F to point to specific
   // address spaces if InferredAddrSpace says so. Postorder is the postorder of
   // all generic address expressions in the use-def graph of function F.
   bool
   rewriteWithNewAddressSpaces(const std::vector<Value *> &Postorder,
                               const ValueToAddrSpaceMapTy &InferredAddrSpace,
                               Function *F) const;
 
else if (auto *MI = dyn_cast<MemIntrinsic>(&I)) {
       // For memset/memcpy/memmove, any pointer operand can be replaced.
       pushPtrOperand(MI->getRawDest());
 
       // Handle 2nd operand for memcpy/memmove.
       if (auto *MTI = dyn_cast<MemTransferInst>(MI))
         pushPtrOperand(MTI->getRawSource());
     } else if (auto *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()) {
+        pushPtrOperand(Cmp->getOperand(0));
+        pushPtrOperand(Cmp->getOperand(1));
+      }
+    }
   }
 
   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) {
       Postorder.push_back(PostorderStack.back().first);
if (auto *MSI = dyn_cast<MemSetInst>(MI)) {
     }
   } else
     llvm_unreachable("unhandled MemIntrinsic");
 
   MI->eraseFromParent();
   return true;
 }
 
+// \p returns true if it is OK to change the address space of constant \p C with
+// a ConstantExpr addrspacecast.
+bool InferAddressSpaces::isSafeToCastConstAddrSpace(Constant *C, unsigned NewAS) const {
+  if (C->getType()->getPointerAddressSpace() == NewAS)
+    return true;
+
+  if (isa<UndefValue>(C) || isa<ConstantPointerNull>(C))
+    return true;
+
+  if (auto *Op = dyn_cast<Operator>(C)) {
+    // If we already have a constant addrspacecast, it should be safe to cast it
+    // off.
+    if (Op->getOpcode() == Instruction::AddrSpaceCast)
+      return isSafeToCastConstAddrSpace(cast<Constant>(Op->getOperand(0)), NewAS);
+
+    if (Op->getOpcode() == Instruction::IntToPtr &&
+        Op->getType()->getPointerAddressSpace() == FlatAddrSpace)
+      return true;
+  }
+
+  return false;
+}
+
 static Value::use_iterator skipToNextUser(Value::use_iterator I,
                                           Value::use_iterator E) {
   User *CurUser = I->getUser();
   ++I;
 
   while (I != E && I->getUser() == CurUser)
     ++I;
 
for (I = V->use_begin(), E = V->use_end(); I != E; ) {
       }
 
       if (auto *II = dyn_cast<IntrinsicInst>(CurUser)) {
         if (rewriteIntrinsicOperands(II, V, NewV))
           continue;
       }
 
       if (isa<Instruction>(CurUser)) {
+        if (ICmpInst *Cmp = dyn_cast<ICmpInst>(CurUser)) {
+          // 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
+
+          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;
+            }
+          }
+
+          // Even if the type mismatches, we can cast the constant.
+          if (auto *KOtherSrc = dyn_cast<Constant>(OtherSrc)) {
+            if (isSafeToCastConstAddrSpace(KOtherSrc, NewAS)) {
+              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: %cmp = icmp eq i32 addrspace(3)* %group.ptr.0, addrspacecast (i32 addrspace(4)* null to i32 addrspace(3)*)
+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_group_flat_cmp_constant_inttoptr(
+; CHECK: %cmp = icmp eq i32 addrspace(3)* %group.ptr.0, addrspacecast (i32 addrspace(4)* inttoptr (i64 400 to i32 addrspace(4)*) to i32 addrspace(3)*)
+define i1 @icmp_group_flat_cmp_constant_inttoptr(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, inttoptr (i64 400 to i32 addrspace(4)*)
+  ret i1 %cmp
+}
+
+; CHECK-LABEL: @icmp_mismatch_flat_group_private_cmp_null(
+; CHECK: %cmp = icmp eq i32 addrspace(3)* %group.ptr.0, addrspacecast (i32* null to i32 addrspace(3)*)
+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
+}
+
+@lds0 = internal addrspace(3) global i32 0, align 4
+@global0 = internal addrspace(1) global i32 0, align 4
+
+; CHECK-LABEL: @icmp_mismatch_flat_group_global_cmp_gv(
+; CHECK: %1 = addrspacecast i32 addrspace(3)* %group.ptr.0 to i32 addrspace(4)*
+; CHECK: %cmp = icmp eq i32 addrspace(4)* %1, addrspacecast (i32 addrspace(1)* @global0 to i32 addrspace(4)*)
+define i1 @icmp_mismatch_flat_group_global_cmp_gv(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 addrspace(1)* @global0 to i32 addrspace(4)*)
+  ret i1 %cmp
+}
+
+; CHECK-LABEL: @icmp_mismatch_group_global_cmp_gv_gv(
+; CHECK: %cmp = icmp eq i32 addrspace(4)* addrspacecast (i32 addrspace(3)* @lds0 to i32 addrspace(4)*), addrspacecast (i32 addrspace(1)* @global0 to i32 addrspace(4)*)
+define i1 @icmp_mismatch_group_global_cmp_gv_gv(i32 addrspace(3)* %group.ptr.0) #0 {
+  %cmp = icmp eq i32 addrspace(4)* addrspacecast (i32 addrspace(3)* @lds0 to i32 addrspace(4)*), addrspacecast (i32 addrspace(1)* @global0 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
+}
+
+; Test non-canonical orders
+; CHECK-LABEL: @icmp_mismatch_flat_group_private_cmp_null_swap(
+; CHECK: %cmp = icmp eq i32 addrspace(3)* addrspacecast (i32* null to i32 addrspace(3)*), %group.ptr.0
+define i1 @icmp_mismatch_flat_group_private_cmp_null_swap(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)* addrspacecast (i32* null to i32 addrspace(4)*), %cast0
+  ret i1 %cmp
+}
+
+; CHECK-LABEL: @icmp_group_flat_cmp_undef_swap(
+; CHECK: %cmp = icmp eq i32 addrspace(3)* undef, %group.ptr.0
+define i1 @icmp_group_flat_cmp_undef_swap(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)* undef, %cast0
+  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