Support mask analysis against 0, and between two scalars

[kile01]

Support was added for two new scenarios:
arith.cmpi ge %scalar, %c0: aka offset comparison to the lower bound of 0. Mask analysis already has an implicit assumption that the beginning of a mask starts at 0, so support was added to allow this case through and assumes that this comparison evaluates to true.
arith.cmpi slt %scalar,_1 %scalar_2: offset comparison between two scalars. E.g.:

%11 = tt.expand_dims %offset
%cst_4 = arith.constant dense<324> : tensor<16x1xi64>
%23 = arith.cmpi slt, %11, %cst_4 : tensor<16x1xi64>

This example is notable in that we cannot take the normal approach of computing the minimum of the lhs and rhs as the new dimension (the lhs offset may be 0). To handle this, a ternary operator is inserted to evaluate the comparison at runtime. If it succeeds, we keep the existing dimensions from the lhs, otherwise we assume nothing should be loaded/stored.

This change also adds a dump method to both MaskState and PtrState as a small QOL improvement.

[kile01]

@microsoft-github-policy-service agree company="Meta"

[nhat-nguyen]

@kile01 Thank you Kirsten! Sorry for the delay as I went on vacation. Looks like there's a typo in the lit test command. The command we use is:

RUN: triton-shared-opt ... %s | FileCheck %s

[kile01]

triton-shared-opt

No worries, this isn't urgent :) and thank you for the catch!

[nhat-nguyen]

Changes look good to me! This is optional, but would be great if we could also have a small snippet of the triton program that generates the IR (I usually put this as a comment in the lit test)

[kile01]

thanks for the review @nhat-nguyen! I've added the example program as a comment

[kile01]

@nhat-nguyen would you mind doing the honors and merging the PR? I don't seem to have write access

stale