Revert "Fix AxisInfo handling of PoisonOp producing MemDesc (#8489)"

This reverts commit 4734af3.

Author: whitneywhtsang

lib/Analysis/AxisInfo.cpp

@@ -276,7 +276,7 @@ class PoisonOpAxisInfoVisitor final : public AxisInfoVisitorImpl<ub::PoisonOp> {
   getAxisInfo(ub::PoisonOp op,
               ArrayRef<const dataflow::Lattice<AxisInfo> *> operands) override {
     unsigned rank = 1;
-    if (auto shape = dyn_cast<RankedTensorType>(op.getType()))
+    if (auto shape = dyn_cast<mlir::ShapedType>(op.getType()))
       rank = shape.getRank();
 
     // Poison values are never accessed, thus assume optimistic values.
@@ -1227,7 +1227,6 @@ void AxisInfo::initDimVectorFromHint(Attribute attr, DimVectorT *vec) {
     return rhs;
   if (rhs.getRank() == 0)
     return lhs;
-  assert(lhs.getRank() == rhs.getRank() && "Mismatched ranks");
   DimVectorT contiguity;
   DimVectorT divisibility;
   DimVectorT constancy;

test/TritonGPU/pipeline-assign-latencies.mlir

@@ -1147,21 +1147,3 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.targ
     tt.return
   }
 }
-
-// -----
-
-// Test that ub.poison producing a memdesc does not get treated like a tensor
-// value in AxisInfo analysis.
-#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 16}>
-#smem = #ttg.shared_memory
-module attributes {"ttg.num-warps" = 4 : i32} {
-  tt.func public @minimal_crash(%lb: i32, %ub: i32) -> !ttg.memdesc<2x2xf16, #shared, #smem, mutable> {
-    %c1 = arith.constant 1 : i32
-    %poison = ub.poison : !ttg.memdesc<2x2xf16, #shared, #smem, mutable>
-    %normal = ttg.local_alloc : () -> !ttg.memdesc<2x2xf16, #shared, #smem, mutable>
-    %result = scf.for %i = %lb to %ub step %c1 iter_args(%current = %poison) -> !ttg.memdesc<2x2xf16, #shared, #smem, mutable> : i32 {
-      scf.yield %normal : !ttg.memdesc<2x2xf16, #shared, #smem, mutable>
-    }
-    tt.return %result : !ttg.memdesc<2x2xf16, #shared, #smem, mutable>
-  }
-}