TEST

Signed-off-by: Ettore Tiotto <[email protected]>

Author: etiotto

lib/Analysis/AxisInfo.cpp

@@ -14,6 +14,8 @@
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
 
+// #define OLD 1
+
 namespace mlir::triton {
 namespace {
 
@@ -1083,7 +1085,12 @@ LogicalResult AxisInfoAnalysis::visitOperation(
   // but why is scf.if not initialized otherwise?
   for (auto op : operands)
     if (op->getValue().getRank() == 0)
+#ifdef OLD
       setToEntryState((dataflow::Lattice<AxisInfo> *)op);
+#else
+      return success();
+#endif
+
   AxisInfo curr = visitors.apply(op, operands);
   if (curr.getRank() == 0) {
     setAllToEntryStates(results);
@@ -1112,9 +1119,17 @@ void AxisInfoAnalysis::visitForOpInductionVar(
   ProgramPoint *programPoint = getProgramPointAfter(op);
   auto *lbLattice = getLatticeElementFor(programPoint, op.getLowerBound());
   auto *stepLattice = getLatticeElementFor(programPoint, op.getStep());
+#ifdef OLD
   for (auto op_iter : {lbLattice, stepLattice})
     if (op_iter->getValue().getRank() == 0)
       setToEntryState((dataflow::Lattice<AxisInfo> *)op_iter);
+#else
+  // If lb or step is not yet ready, skip this operation for now.
+  if (lbLattice->getValue().getRank() == 0 ||
+      stepLattice->getValue().getRank() == 0) {
+    return;
+  }
+#endif
 
   AxisInfo::DimVectorT knownContiguity(1, 1);
   AxisInfo::DimVectorT knownDivisibility(1, 1);
@@ -1188,16 +1203,25 @@ void AxisInfo::initDimVectorFromHint(Attribute attr, DimVectorT *vec) {
       initPessimisticStateFromFunc(blockArg.getArgNumber(), fun,
                                    &knownContiguity, &knownDivisibility,
                                    &knownConstancy);
-    } else if (isa<RegionBranchOpInterface, gpu::WarpSpecializePartitionsOp>(
-                   op)) {
+    }
+#ifdef OLD
+    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.
+#else
+    else if (isa<gpu::WarpSpecializePartitionsOp>(op)) {
+      // Initialize the arguments to gpu::WarpSpecializePartitionsOp with
+      // "unknown" state: the maximum possible divisibility, contiguity, and
+      // constancy.
+#endif
       knownDivisibility = DimVectorT(rank, kMaxDivisor);
       knownConstancy = DimVectorT(rank, kMaxDivisor);
       knownContiguity = DimVectorT(rank, kMaxDivisor);
     }
   } else if (Operation *op = value.getDefiningOp()) {
+#ifdef OLD
     if (isa<RegionBranchOpInterface>(op)) {
       // scf::ForOp, scf::IfOp, scf::WhileOp
       // Control flow operations are initialized with "unknown" state:
@@ -1206,6 +1230,7 @@ void AxisInfo::initDimVectorFromHint(Attribute attr, DimVectorT *vec) {
       knownConstancy = DimVectorT(rank, kMaxDivisor);
       knownContiguity = DimVectorT(rank, kMaxDivisor);
     }
+#endif
     // Other operations are conservatively initialized with the lowest possible
     // divisibility, contiguity, and constancy unless they have specified.
     AxisInfo::initDimVectorFromHint(op->getDiscardableAttr("tt.divisibility"),
@@ -1358,6 +1383,12 @@ void ModuleAxisInfoAnalysis::initialize(FunctionOpInterface funcOp,
   auto *axisInfoMap = getFuncData(funcOp);
   auto updateAxisInfoMap = [&](Value value) {
     auto axisInfo = analysis->getLatticeElement(value)->getValue();
+#ifdef OLD
+    // If we could not determine the AxisInfo for this value, assume the
+    // pessimistic state.
+    if (axisInfo.getRank() == 0)
+      axisInfo = AxisInfo::getPessimisticValueState(value);
+#endif
     auto &valInfo = (*axisInfoMap)[value];
     valInfo = AxisInfo::join(axisInfo, valInfo);
   };