Handle mix of structured and unstructured accesses in loops

[nhat-nguyen]

Previously, before the conversion to the PtrDialect for fallback, the structured-to-memref pass has to convert loop's iter-args with triton pointer to unranked memref. This conversion ensures all types coming out of triton-shared are mlir built-in types and therefore allows the CPU backend to correctly lower the IR to llvm. However, in reality, structured ops do not need to use the loop iter-args since ptr-analysis generates load/store ops that directly use the kernel arguments as source; this means the conversion is mostly unnecessary.

With the introduction of the fallback using the PtrDialect (triton-to-ptr), we also convert the loop iter-args of triton pointer type to PtrDialect's ptr type. This conversion, along with the conversion to unranked memref above, means we will end up with unrealized_conversion_cast ops that convert back and forth between these two types when handling triton programs that have mixed uses of structured and unstructured accesses in loops.

To solve this issue, we:

• remove the conversion of loop-iter of triton ptr type to unranked memref since it is unnecessary as described as above
• run remove-dead-values to remove unused loop-iter args; this pass previously could not run in presence of ops with arbitrary regions but has now been fixed in this PR: [mlir] [liveness] Conservatively mark operands of return-like op inside non-callable and non-regionbranch op as live llvm/llvm-project#140793. Running remove-dead-values gives two benefits:
  • ability to remove all unused loop iter-arg that isn't used after ptr-analysis
  • make our codegen more efficient in general