[llvm] e2f6290 - [VectorCombine] Discard ScalarizationResult state in early exit.

[Florian Hahn via llvm-commits]

Author: Florian Hahn
Date: 2021-09-28T12:52:16+01:00
New Revision: e2f6290e06be63149af770197b772248369ac038

URL: https://github.com/llvm/llvm-project/commit/e2f6290e06be63149af770197b772248369ac038
DIFF: https://github.com/llvm/llvm-project/commit/e2f6290e06be63149af770197b772248369ac038.diff

LOG: [VectorCombine] Discard ScalarizationResult state in early exit.

ScalarizationResult's destructor makes sure ToFreeze is not ignored if
set. Currently, scalarizeLoadExtract has an early exit if the index is
not safe directly. But when it is SafeWithFreeze, we need to discard the
state first, otherwise we hit the assert in the destructor.

Fixes PR51992.

Added: 
    

Modified: 
    llvm/lib/Transforms/Vectorize/VectorCombine.cpp
    llvm/test/Transforms/VectorCombine/AArch64/load-extractelement-scalarization.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
index dcf7b68bd7923..2038655051b37 100644
--- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
+++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
@@ -822,6 +822,12 @@ class ScalarizationResult {
   /// freeze.
   bool isSafeWithFreeze() const { return Status == StatusTy::SafeWithFreeze; }
 
+  /// Reset the state of Unsafe and clear ToFreze if set.
+  void discard() {
+    ToFreeze = nullptr;
+    Status = StatusTy::Unsafe;
+  }
+
   /// Freeze the ToFreeze and update the use in \p User to use it.
   void freeze(IRBuilder<> &Builder, Instruction &UserI) {
     assert(isSafeWithFreeze() &&
@@ -1006,6 +1012,7 @@ bool VectorCombine::scalarizeLoadExtract(Instruction &I) {
     auto ScalarIdx = canScalarizeAccess(FixedVT, UI->getOperand(1), &I, AC, DT);
     if (!ScalarIdx.isSafe()) {
       // TODO: Freeze index if it is safe to do so.
+      ScalarIdx.discard();
       return false;
     }
 

diff  --git a/llvm/test/Transforms/VectorCombine/AArch64/load-extractelement-scalarization.ll b/llvm/test/Transforms/VectorCombine/AArch64/load-extractelement-scalarization.ll
index e14929c35318e..3a9fa52905c80 100644
--- a/llvm/test/Transforms/VectorCombine/AArch64/load-extractelement-scalarization.ll
+++ b/llvm/test/Transforms/VectorCombine/AArch64/load-extractelement-scalarization.ll
@@ -678,3 +678,37 @@ define i32 @load_multiple_extracts_with_variable_indices_large_vector_all_valid_
   ret i32 %res
 }
 
+; Test case from PR51992.
+; TODO: could handle by inserting freeze.
+define i8 @load_extract_safe_with_freeze(<8 x i8> %in, <16 x i8>* %src) {
+; CHECK-LABEL: @load_extract_safe_with_freeze(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[EXT_IDX:%.*]] = extractelement <8 x i8> [[IN:%.*]], i32 0
+; CHECK-NEXT:    [[EXT_IDX_I32:%.*]] = zext i8 [[EXT_IDX]] to i32
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[EXT_IDX]], 99
+; CHECK-NEXT:    br i1 [[CMP]], label [[THEN:%.*]], label [[EXIT:%.*]]
+; CHECK:       then:
+; CHECK-NEXT:    [[LOAD:%.*]] = load <16 x i8>, <16 x i8>* [[SRC:%.*]], align 16
+; CHECK-NEXT:    [[AND:%.*]] = and i32 [[EXT_IDX_I32]], 15
+; CHECK-NEXT:    [[EXT:%.*]] = extractelement <16 x i8> [[LOAD]], i32 [[AND]]
+; CHECK-NEXT:    br label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[P:%.*]] = phi i8 [ 0, [[ENTRY:%.*]] ], [ [[EXT]], [[THEN]] ]
+; CHECK-NEXT:    ret i8 0
+;
+entry:
+  %ext.idx = extractelement <8 x i8> %in, i32 0
+  %ext.idx.i32 = zext i8 %ext.idx to i32
+  %cmp = icmp ult i8 %ext.idx, 99
+  br i1 %cmp, label %then, label %exit
+
+then:
+  %load = load <16 x i8>, <16 x i8>* %src, align 16
+  %and = and i32 %ext.idx.i32, 15
+  %ext = extractelement <16 x i8> %load, i32 %and
+  br label %exit
+
+exit:
+  %p = phi i8 [ 0, %entry ], [ %ext, %then ]
+  ret i8 0
+}