Revert "[Reland] Fix handling of unvisited operands in AxisInfoAnalysis (#8758)"

This reverts commit 31281bc.

Author: anmyachev

lib/Analysis/AxisInfo.cpp

@@ -1079,10 +1079,11 @@ AxisInfoAnalysis::AxisInfoAnalysis(DataFlowSolver &solver,
 LogicalResult AxisInfoAnalysis::visitOperation(
     Operation *op, ArrayRef<const dataflow::Lattice<AxisInfo> *> operands,
     ArrayRef<dataflow::Lattice<AxisInfo> *> results) {
-  // If any operands are not yet ready, skip this operation for now.
+  // TODO: For sure not the right way to do this
+  // but why is scf.if not initialized otherwise?
   for (auto op : operands)
     if (op->getValue().getRank() == 0)
-      return success();
+      setToEntryState((dataflow::Lattice<AxisInfo> *)op);
   AxisInfo curr = visitors.apply(op, operands);
   if (curr.getRank() == 0) {
     setAllToEntryStates(results);
@@ -1111,11 +1112,9 @@ void AxisInfoAnalysis::visitForOpInductionVar(
   ProgramPoint *programPoint = getProgramPointAfter(op);
   auto *lbLattice = getLatticeElementFor(programPoint, op.getLowerBound());
   auto *stepLattice = getLatticeElementFor(programPoint, op.getStep());
-  // If lb or step is not yet ready, skip this operation for now.
-  if (lbLattice->getValue().getRank() == 0 ||
-      stepLattice->getValue().getRank() == 0) {
-    return;
-  }
+  for (auto op_iter : {lbLattice, stepLattice})
+    if (op_iter->getValue().getRank() == 0)
+      setToEntryState((dataflow::Lattice<AxisInfo> *)op_iter);
 
   AxisInfo::DimVectorT knownContiguity(1, 1);
   AxisInfo::DimVectorT knownDivisibility(1, 1);
@@ -1189,15 +1188,24 @@ void AxisInfo::initDimVectorFromHint(Attribute attr, DimVectorT *vec) {
       initPessimisticStateFromFunc(blockArg.getArgNumber(), fun,
                                    &knownContiguity, &knownDivisibility,
                                    &knownConstancy);
-    } else if (isa<gpu::WarpSpecializePartitionsOp>(op)) {
-      // Initialize the arguments to gpu::WarpSpecializePartitionsOp with
-      // "unknown" state: the maximum possible divisibility, contiguity, and
-      // constancy.
+    } else if (isa<RegionBranchOpInterface, gpu::WarpSpecializePartitionsOp>(
+                   op)) {
+      // scf::ForOp, scf::IfOp, scf::WhileOp, gpu::WarpSpecializePartitionsOp
+      // Control flow operations are initialized with "unknown" state:
+      // the maximum possible divisibility, contiguity, and constancy.
       knownDivisibility = DimVectorT(rank, kMaxDivisor);
       knownConstancy = DimVectorT(rank, kMaxDivisor);
       knownContiguity = DimVectorT(rank, kMaxDivisor);
     }
   } else if (Operation *op = value.getDefiningOp()) {
+    if (isa<RegionBranchOpInterface>(op)) {
+      // scf::ForOp, scf::IfOp, scf::WhileOp
+      // Control flow operations are initialized with "unknown" state:
+      // the maximum possible divisibility, contiguity, and constancy.
+      knownDivisibility = DimVectorT(rank, kMaxDivisor);
+      knownConstancy = DimVectorT(rank, kMaxDivisor);
+      knownContiguity = DimVectorT(rank, kMaxDivisor);
+    }
     // Other operations are conservatively initialized with the lowest possible
     // divisibility, contiguity, and constancy unless they have specified.
     AxisInfo::initDimVectorFromHint(op->getDiscardableAttr("tt.divisibility"),
@@ -1350,10 +1358,6 @@ void ModuleAxisInfoAnalysis::initialize(FunctionOpInterface funcOp,
   auto *axisInfoMap = getFuncData(funcOp);
   auto updateAxisInfoMap = [&](Value value) {
     auto axisInfo = analysis->getLatticeElement(value)->getValue();
-    // If we could not determine the AxisInfo for this value, assume the
-    // pessimistic state.
-    if (axisInfo.getRank() == 0)
-      axisInfo = AxisInfo::getPessimisticValueState(value);
     AxisInfo curAxisInfo;
     if (axisInfoMap->count(value)) {
       curAxisInfo = AxisInfo::join(axisInfo, axisInfoMap->lookup(value));

test/Analysis/test-alignment.mlir

@@ -1089,18 +1089,3 @@ tt.func public @test_inductor_for() {
   }
   tt.return
 }
-
-// -----
-
-// Verify that if an operation is statically determined to be dead, we fall back
-// to assigning it a pessimistic value, rather than skipping it entirely.
-tt.func @dead_op_pessimistic() {
-  %c5 = arith.constant dense<5> : tensor<4xi32>
-  %c7 = arith.constant dense<7> : tensor<4xi32>
-  %false = arith.constant false
-  scf.if %false {
-    // expected-remark @below {{contiguity = [1], divisibility = [1], constancy = [1], constant_value = <none>}}
-    %add = arith.addi %c5, %c7 : tensor<4xi32>
-  }
-  tt.return
-}