[BACKEND] Support of ConvertLayoutOp from blocked to blocked and SliceLayout with blocked parent (#658)

lib/Analysis/AxisInfo.cpp

@@ -1,4 +1,5 @@
 #include "mlir/Analysis/DataFlowAnalysis.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "llvm/Support/raw_ostream.h"
 #include <iostream>
 
@@ -46,6 +47,11 @@ AxisInfo AxisInfo::getPessimisticValueState(Value value) {
           fun.getArgAttr(blockArg.getArgNumber(), "tt.divisibility");
       if (attr)
         divHint = attr.cast<IntegerAttr>().getValue().getZExtValue();
+    } else if (auto fun = dyn_cast<LLVM::LLVMFuncOp>(op)) {
+      Attribute attr =
+          fun.getArgAttr(blockArg.getArgNumber(), "tt.divisibility");
+      if (attr)
+        divHint = attr.cast<IntegerAttr>().getValue().getZExtValue();
     }
   }
   DimVectorT contiguity(rank, 1);
@@ -203,6 +209,13 @@ ChangeResult AxisInfoAnalysis::visitOperation(
     }
     curr = AxisInfo(contiguity, divisibility, constancy);
   }
+  // UnrealizedConversionCast
+  // This is needed by TritonGPUToLLVM, to get AxisInfo when the graph is
+  // in the process of a PartialConversion, where UnrealizedConversionCast
+  // may exist
+  if (llvm::isa<mlir::UnrealizedConversionCastOp>(op)) {
+    curr = operands[0]->getValue();
+  }
   if (curr.getRank() == 0) {
     return markAllPessimisticFixpoint(op->getResults());
   }