[InstCombine] Promote expression tree with @llvm.vscale when zero-extending result.

[sdesmalen]

Summary

The LoopVectorizer emits the (scaled) element count as i32, which for
scalable VFs results in calls to @llvm.vscale.i32(). This value is scaled
and further zero-extended to i64.

The zero-extend can be folded away by executing the whole expression in i64
type using @llvm.vscale.i64(). Any logical and that would needed to mask
the result can be further folded away by KnownBits analysis when
vscale_range is set.

[goldstein.w.n]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

260

Since Res is conditionally set here (i.e we can have Call w.o it being vscale think this may segfault b.c Res is set nullptr above and unconditionally derefenced below.

llvm/test/Transforms/InstCombine/AArch64/vscale.ll

1 ↗

(On Diff #493729)

Can the tests the split into a separate patch so we can see the diff?

[fhahn]

llvm/test/Transforms/InstCombine/AArch64/vscale.ll

1 ↗

(On Diff #493729)

Does the test need to be AArch64 specific? IIUC this is a target-independent fold of @llvm.vscale, so it might be better to have a generic test?

[sdesmalen]

• Fixed issue where Res was conditionally set but still used.
• Made test agnostic of target.
• Patch now only shows the diff for the test.

[sdesmalen]

Thanks for the quick feedback!

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

260

Good spot!

[goldstein.w.n]

llvm/test/Transforms/InstCombine/vscale.ll

18

Is this semantically equivilent? Unfortunately alive2 doesn't support @llvm.vscale but is it possible for anything in vscale[29:31] to be set? If so the result will be different.

[craig.topper]

llvm/test/Transforms/InstCombine/vscale.ll

18

The vscale_range attribute is placing a limit on the upper bound. I think without that there would be an AND instruction at the end to clear the upper bits.

[sdesmalen]

• Added a test without vscale_range, to show it results in explicit anding of upper bits.
• Replaced LLVM_FALLTHROUGH with explicit 'llvm_unreachable' and added a switch-statement for the intrinsic ID.

[nikic]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

case Instruction::PHI: {
     Res = NPN;
     break;
   }
   case Instruction::FPToUI:
   case Instruction::FPToSI:
     Res = CastInst::Create(
       static_cast<Instruction::CastOps>(Opc), I->getOperand(0), Ty);
     break;
+  case Instruction::Call:
+    if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
+      switch (II->getIntrinsicID()) {
+      default:
+        llvm_unreachable("Unsupported call!");
+      case Intrinsic::vscale: {
+        Function *Fn =
+            Intrinsic::getDeclaration(I->getModule(), Intrinsic::vscale, {Ty});
+        Res = CallInst::Create(Fn->getFunctionType(), Fn);
+        break;
+      }
+      }
+    }
+    break;
   default:
     // TODO: Can handle more cases here.
     llvm_unreachable("Unreachable!");
   }
 
   Res->takeName(I);
   return InsertNewInstWith(Res, *I);
 }
case Instruction::PHI: {
     for (unsigned i = 1, e = PN->getNumIncomingValues(); i != e; ++i)
       if (!canEvaluateZExtd(PN->getIncomingValue(i), Ty, Tmp, IC, CxtI) ||
           // TODO: If important, we could handle the case when the BitsToClear
           // are known zero in the disagreeing input.
           Tmp != BitsToClear)
         return false;
     return true;
   }
+  case Instruction::Call:
+    // llvm.vscale() can always be executed in larger type, because the
+    // value is automatically zero-extended.
+    if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(I))
+      if (II->getIntrinsicID() == Intrinsic::vscale)
+        return true;
+    return false;
   default:
     // TODO: Can handle more cases here.
     return false;
   }
 }
 
 Instruction *InstCombinerImpl::visitZExt(ZExtInst &Zext) {
   // If this zero extend is only used by a truncate, let the truncate be

llvm/test/Transforms/InstCombine/vscale.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes='instcombine' -S < %s | FileCheck %s
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+
+define i64 @promote_vscale_i32_to_i64() {
+; CHECK-LABEL: @promote_vscale_i32_to_i64(
+; CHECK-NEXT:    [[VSCALE:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    [[EXT:%.*]] = and i64 [[VSCALE]], 4294967295
+; CHECK-NEXT:    ret i64 [[EXT]]
+;
+  %vscale = call i32 @llvm.vscale.i32()
+  %ext = zext i32 %vscale to i64
+  ret i64 %ext
+}
+
+define i64 @pomote_zext_shl_vscale_i32_to_i64() {
+; CHECK-LABEL: @pomote_zext_shl_vscale_i32_to_i64(
+; CHECK-NEXT:    [[VSCALE:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    [[SHL:%.*]] = shl i64 [[VSCALE]], 3
+; CHECK-NEXT:    [[EXT:%.*]] = and i64 [[SHL]], 4294967288
+; CHECK-NEXT:    ret i64 [[EXT]]
+;
+  %vscale = call i32 @llvm.vscale.i32()
+  %shl = shl i32 %vscale, 3
+  %ext = zext i32 %shl to i64
+  ret i64 %ext
+}
+
+; Same test as @pomote_zext_shl_vscale_i32_to_i64, but with the
+; vscale_range attribute so that the 'and' is folded away.
+define i64 @free_zext_vscale_shl_i32_to_i64() #0 {
+; CHECK-LABEL: @free_zext_vscale_shl_i32_to_i64(
+; CHECK-NEXT:    [[VSCALE:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    [[SHL:%.*]] = shl nuw nsw i64 [[VSCALE]], 3
+; CHECK-NEXT:    ret i64 [[SHL]]
+;
+  %vscale = call i32 @llvm.vscale.i32()
+  %shl = shl i32 %vscale, 3
+  %ext = zext i32 %shl to i64
+  ret i64 %ext
+}
+
+declare i32 @llvm.vscale.i32()
+
+attributes #0 = { vscale_range(1,16) }