[SVE] Add folds for sign and zero extends of vscale

[Matt]

llvm/test/Transforms/InstCombine/vscale_sext_and_zext.ll

44

Nit: Missing space.

[efriedma]

I think you might need to be a bit more careful to justify these transforms.

LangRef says "If the result value does not fit in the result type, then the result is a poison value." But does that mean the sign bit is guaranteed to be zero?

[DylanFleming-arm]

We could use the max-vscale attribute to determine if the runtime value fits the type. If this cannot be determined, the optimization would be skipped, otherwise we change the type of llvm.vscale to the extended type.
Would that be a suitable approach?

[efriedma]

I assume the max-vscale attribute will usually be present? In that case, checking it seems best.

(We might actually want to consider relaxing the LangRef rule; if we can easily determine whether the result fits into a given bitwidth, there isn't really any reason to make llvm.vscale return poison.)

[DylanFleming-arm]

There's instances where max-vscale won't be known at compile time, but this route should allow for the optimizations wherever possible.

I believe in terms of the LangRef, it's still possible for llvm.vscale to return poision, and it could still happen in cases where max-vscale is unknown until runtime. So personally I don't think any adjustment is required.

[sdesmalen]

In D105994#2881236, @efriedma wrote:

I assume the max-vscale attribute will usually be present? In that case, checking it seems best.

(We might actually want to consider relaxing the LangRef rule; if we can easily determine whether the result fits into a given bitwidth, there isn't really any reason to make llvm.vscale return poison.)

Just to make sure that I understand the suggestion, do you mean changing the wording in the LangRef that the result is undefined instead of poison? If so, does that mean we could replace llvm.vscale.i8() with undef if max(vscale_range) doesn't fit the i8?

[efriedma]

In D105994#2890403, @sdesmalen wrote:

In D105994#2881236, @efriedma wrote:

I assume the max-vscale attribute will usually be present? In that case, checking it seems best.

(We might actually want to consider relaxing the LangRef rule; if we can easily determine whether the result fits into a given bitwidth, there isn't really any reason to make llvm.vscale return poison.)

Just to make sure that I understand the suggestion, do you mean changing the wording in the LangRef that the result is undefined instead of poison? If so, does that mean we could replace llvm.vscale.i8() with undef if max(vscale_range) doesn't fit the i8?

My suggestion is that, for example, @llvm.vscale.i8() returns "vscale & 0xFF". Not really sure if it's worthwhile, but it seems like an unnecessary source of poison.

[sdesmalen]

In D105994#2890694, @efriedma wrote:

Just to make sure that I understand the suggestion, do you mean changing the wording in the LangRef that the result is undefined instead of poison? If so, does that mean we could replace llvm.vscale.i8() with undef if max(vscale_range) doesn't fit the i8?

My suggestion is that, for example, @llvm.vscale.i8() returns "vscale & 0xFF". Not really sure if it's worthwhile, but it seems like an unnecessary source of poison.

What problem would that solve?

[efriedma]

In D105994#2890907, @sdesmalen wrote:

In D105994#2890694, @efriedma wrote:

Just to make sure that I understand the suggestion, do you mean changing the wording in the LangRef that the result is undefined instead of poison? If so, does that mean we could replace llvm.vscale.i8() with undef if max(vscale_range) doesn't fit the i8?

My suggestion is that, for example, @llvm.vscale.i8() returns "vscale & 0xFF". Not really sure if it's worthwhile, but it seems like an unnecessary source of poison.

What problem would that solve?

It would allow folding "trunc vscale". And it might be a little simpler to understand. Either way, not a big deal. :)

[DylanFleming-arm]

Added check for vscale_range attribute before optimisation
If the attribute isn't present, or if the maximum value exceeds the bitwidth of the original instrinsic, the optimization is skipped

Updated .ll test to test the extra logic

[efriedma]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

1371

This shift can overflow.

I'd suggest using Log2_32(MaxVScale) instead.

[DylanFleming-arm]

Changed bitshift to Log2_32

[efriedma]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

Instruction *InstCombinerImpl::visitZExt(ZExtInst &CI) {
   Value *And;
   if (SrcI && match(SrcI, m_OneUse(m_Xor(m_Value(And), m_Constant(C)))) &&
       match(And, m_OneUse(m_And(m_Trunc(m_Value(X)), m_Specific(C)))) &&
       X->getType() == CI.getType()) {
     Constant *ZC = ConstantExpr::getZExt(C, CI.getType());
     return BinaryOperator::CreateXor(Builder.CreateAnd(X, ZC), ZC);
   }
 
+  if (match(Src, m_VScale(DL))) {
+    if (CI.getFunction()->hasFnAttribute(Attribute::VScaleRange)) {
+      unsigned MaxVScale = CI.getFunction()
+                               ->getFnAttribute(Attribute::VScaleRange)
+                               .getVScaleRangeArgs()
+                               .second;
+      unsigned TypeWidth = Src->getType()->getScalarSizeInBits();
+      if (Log2_32(MaxVScale) < TypeWidth) {
+        Value *VScale = Builder.CreateVScale(ConstantInt::get(DestTy, 1));
+        return replaceInstUsesWith(CI, VScale);
+      }
+    }
+  }
+
   return nullptr;
 }
 
 /// Transform (sext icmp) to bitwise / integer operations to eliminate the icmp.
 Instruction *InstCombinerImpl::transformSExtICmp(ICmpInst *ICI,
                                                  Instruction &CI) {
   Value *Op0 = ICI->getOperand(0), *Op1 = ICI->getOperand(1);
   ICmpInst::Predicate Pred = ICI->getPredicate();
Constant *NewShAmt = ConstantExpr::getSub(
         ConstantInt::get(DestTy, DestTy->getScalarSizeInBits()),
         NumLowbitsLeft);
     NewShAmt =
         Constant::mergeUndefsWith(Constant::mergeUndefsWith(NewShAmt, BA), CA);
     A = Builder.CreateShl(A, NewShAmt, CI.getName());
     return BinaryOperator::CreateAShr(A, NewShAmt);
   }
 
+  if (match(Src, m_VScale(DL))) {
+    if (CI.getFunction()->hasFnAttribute(Attribute::VScaleRange)) {
+      unsigned MaxVScale = CI.getFunction()
+                               ->getFnAttribute(Attribute::VScaleRange)
+                               .getVScaleRangeArgs()
+                               .second;
+      unsigned TypeWidth = Src->getType()->getScalarSizeInBits();
+      if (Log2_32(MaxVScale) < (TypeWidth - 1)) {
+        Value *VScale = Builder.CreateVScale(ConstantInt::get(DestTy, 1));
+        return replaceInstUsesWith(CI, VScale);
+      }
+    }
+  }
+
   return nullptr;
 }
 
 /// Return a Constant* for the specified floating-point constant if it fits
 /// in the specified FP type without changing its value.
 static bool fitsInFPType(ConstantFP *CFP, const fltSemantics &Sem) {
   bool losesInfo;
   APFloat F = CFP->getValueAPF();

llvm/test/Transforms/InstCombine/vscale_sext_and_zext.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -instcombine -S | FileCheck %s
+
+define i64 @vscale_SExt_i32toi64() #0 {
+; CHECK:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    ret i64 [[TMP0]]
+entry:
+  %0 = call i32 @llvm.vscale.i32()
+  %1 = sext i32 %0 to i64
+  ret i64 %1
+}
+
+define i32 @vscale_SExt_i8toi32() #0 {
+; CHECK:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = call i32 @llvm.vscale.i32()
+; CHECK-NEXT:    ret i32 [[TMP0]]
+entry:
+  %0 = call i8 @llvm.vscale.i8()
+  %1 = sext i8 %0 to i32
+  ret i32 %1
+}
+
+
+define i32 @vscale_SExt_i8toi32_poison() vscale_range(0, 192) {
+; CHECK:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = call i8 @llvm.vscale.i8()
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i8 [[TMP0]] to i32
+; CHECK-NEXT:    ret i32 [[TMP1]]
+  entry:
+  %0 = call i8 @llvm.vscale.i8()
+  %1 = sext i8 %0 to i32
+  ret i32 %1
+}
+
+
+
+define i64 @vscale_ZExt_i32toi64() #0 {
+; CHECK:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    ret i64 [[TMP0]]
+entry:
+  %0 = call i32 @llvm.vscale.i32()
+  %1 = zext i32 %0 to i64
+  ret i64 %1
+}
+
+define i64 @vscale_ZExt_i1toi64() vscale_range(0, 1) {
+; CHECK:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = call i64 @llvm.vscale.i64()
+; CHECK-NEXT:    ret i64 [[TMP0]]
+entry:
+  %0 = call i1 @llvm.vscale.i1()
+  %1 = zext i1 %0 to i64
+  ret i64 %1
+}
+
+define i32 @vscale_ZExt_i8toi32_poison() vscale_range(0, 1024) {
+; CHECK:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = call i8 @llvm.vscale.i8()
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[TMP0]] to i32
+; CHECK-NEXT:    ret i32 [[TMP1]]
+  entry:
+  %0 = call i8 @llvm.vscale.i8()
+  %1 = zext i8 %0 to i32
+  ret i32 %1
+}
+
+define i32 @vscale_ZExt_i16toi32_unknown() {
+; CHECK:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = call i16 @llvm.vscale.i16()
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i16 [[TMP0]] to i32
+; CHECK-NEXT:    ret i32 [[TMP1]]
+  entry:
+  %0 = call i16 @llvm.vscale.i16()
+  %1 = zext i16 %0 to i32
+  ret i32 %1
+}
+
+attributes #0 = { vscale_range(0, 16) }
+
+declare i1  @llvm.vscale.i1()
+declare i8  @llvm.vscale.i8()
+declare i16 @llvm.vscale.i16()
+declare i32 @llvm.vscale.i32()