[Analysis] divisibility handling for dividing by a power-of-two constant (#6657)

Given this ttir:
```
tt.func @div(%arg0: i32 {tt.divisibility = 16 : i32}) {
  %1 = arith.constant 2 : i32
  %2 = arith.divsi %arg0, %1 : i32
}
```

Triton should be able to infer that %2 has divisibility = 8.

Use case: a kernel in which I load with a `mask = offsets < numel // 2`.

Author: davidberard98

lib/Analysis/AxisInfo.cpp

@@ -467,6 +467,13 @@ class DivOpAxisInfoVisitor final : public BinaryOpVisitorImpl<OpTy> {
     if (rhs.getConstantValue().has_value() &&
         rhs.getConstantValue().value() == 1)
       return lhs.getDivisibility(dim);
+    // Case 3: lhs has contiguity of 1 in this dimension and rhs is a power of 2
+    if (rhs.getConstantValue().has_value() &&
+        llvm::isPowerOf2_64(std::abs(rhs.getConstantValue().value())) &&
+        lhs.getContiguity(dim) == 1) {
+      int64_t absRhs = std::abs(rhs.getConstantValue().value());
+      return std::max<int64_t>(1, lhs.getDivisibility(dim) / absRhs);
+    }
     // otherwise: return 1
     return 1;
   }

test/Analysis/test-alignment.mlir

@@ -129,7 +129,7 @@ tt.func @mul(%arg0: i64 {tt.divisibility = 16 : i32}) {
 
 // -----
 
-tt.func @div() {
+tt.func @div(%arg0: i32 {tt.divisibility = 16 : i32}) {
   // expected-remark @below {{contiguity = [128], divisibility = [1073741824], constancy = [1], constant_value = <none>}}
   %0 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>
   // expected-remark @below {{contiguity = [1], divisibility = [1], constancy = [128], constant_value = 1}}
@@ -154,6 +154,23 @@ tt.func @div() {
   %10 = tt.make_range {end = 8320 : i32, start = 8192 : i32} : tensor<128xi32>
   // expected-remark @below {{contiguity = [1], divisibility = [1], constancy = [64], constant_value = <none>}}
   %11 = arith.divsi %10, %4 : tensor<128xi32>
+  // expected-remark @below {{contiguity = [1], divisibility = [2], constancy = [1], constant_value = 2}}
+  %12 = arith.constant 2 : i32
+  // dividing a scalar by a power of two should give predictable divisibility
+  // expected-remark @below {{contiguity = [1], divisibility = [8], constancy = [1], constant_value = <none>}}
+  %13 = arith.divsi %arg0, %12 : i32
+  // expected-remark @below {{contiguity = [1], divisibility = [32], constancy = [1], constant_value = 32}}
+  %14 = arith.constant 32 : i32
+  // expected-remark @below {{contiguity = [1], divisibility = [1], constancy = [1], constant_value = <none>}}
+  %15 = arith.divsi %arg0, %14 : i32
+  // expected-remark @below {{contiguity = [1], divisibility = [2], constancy = [1], constant_value = 6}}
+  %16 = arith.constant 6 : i32
+  // expected-remark @below {{contiguity = [1], divisibility = [1], constancy = [1], constant_value = <none>}}
+  %17 = arith.divsi %arg0, %16 : i32
+  // expected-remark @below {{contiguity = [1], divisibility = [2], constancy = [128], constant_value = 2}}
+  %18 = arith.constant dense<2> : tensor<128xi32>
+  // expected-remark @below {{contiguity = [1], divisibility = [1], constancy = [2], constant_value = <none>}}
+  %19 = arith.divsi %0, %18 : tensor<128xi32>
   tt.return
 }