[TritonGPU] Augment FuseNestedLoops to handle dependent inner loop bounds (#8132)

This PR teaches FuseNestedLoops to handle inner loops whose loop bounds
are some (pure) function of the outer loop bounds. FuseNestedLoops
slices the inner loop bound computations into a loop before the fused
loop to compute the total number of fused iterations. Then, inside the
first fused prologue, the inner loop lengths for the current outer loop
iterations are computed. The pass supports a mix of inner loops whose
bounds can be made outer loop invariant and those that are not.

This patch also adds a small hack that pattern matches `tl.assume(ub >
lb)` for the inner loop bounds to allow speculation (i.e. all inner
loops execute at least once).

Author: Mogball

include/triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h

@@ -40,7 +40,10 @@ bool isPureScalarOp(Operation *op);
 bool getDominatingValueSetOpsToHoist(
     DominanceInfo &domInfo, Operation *refOp, ArrayRef<Value> valueSet,
     llvm::SetVector<Operation *> &toHoist,
-    function_ref<bool(Operation *)> canHoist = isPureScalarOp);
+    function_ref<bool(Operation *)> canHoist = isPureScalarOp,
+    function_ref<bool(BlockArgument)> canUseArg = [](BlockArgument) {
+      return false;
+    });
 
 // Hoist the given set of operations above the reference operation.
 void hoistOpsBefore(Operation *refOp,