Fix MaskAnalysis bool type workaround (#346)

In MaskAnalysis, a Boolean type shouldn't be passed around as Index.
This replaces the workaround added in PR
#318.
We use the following two tests as baseline, they needed the workaround
to pass before and now passing with our fix.
```
test/Conversion/TritonToStructured/mask_ld_st_scalar_dim.mlir
python/examples/test_mask.py
```

Co-authored-by: Linxiao Wang <[email protected]>

Author: lxwangruc

include/triton-shared/Analysis/OpFoldResultUtils.h

@@ -27,6 +27,9 @@ std::optional<int64_t> getIntAttr(const OpFoldResult ofr);
 // attribute or a constant value.
 bool hasConstZero(const OpFoldResult ofr);
 
+// Cast OpFoldResult to Value.
+Value ofrToValue(const OpFoldResult ofr, const Location loc, OpBuilder &b);
+
 // Create a value of index type if necessary from an OpFoldResult.
 Value ofrToIndexValue(const OpFoldResult ofr, const Location loc, OpBuilder &b);
 

lib/Analysis/MaskAnalysis.cpp

@@ -302,9 +302,13 @@ LogicalResult MaskState::parseConstant(arith::ConstantOp constOp,
 LogicalResult MaskState::parseIntScalar(Value scalar, const Location loc,
                                         OpBuilder &builder) {
   assert(this->isEmpty());
-  auto castOp =
-      builder.create<arith::IndexCastOp>(loc, builder.getIndexType(), scalar);
-  this->scalar = castOp.getResult();
+  if (scalar.getType().isInteger(1)) {
+    this->scalar = scalar;
+  } else {
+    auto castOp =
+        builder.create<arith::IndexCastOp>(loc, builder.getIndexType(), scalar);
+    this->scalar = castOp.getResult();
+  }
   return success();
 }
 

lib/Analysis/OpFoldResultUtils.cpp

@@ -52,6 +52,16 @@ bool hasConstZero(const OpFoldResult ofr) {
   return false;
 }
 
+Value ofrToValue(const OpFoldResult ofr, const Location loc, OpBuilder &b) {
+  if (Value val = dyn_cast<Value>(ofr)) {
+    return val;
+  }
+
+  auto attr = dyn_cast<Attribute>(ofr);
+  auto typedAttr = dyn_cast<TypedAttr>(attr);
+  return b.create<arith::ConstantOp>(loc, typedAttr);
+}
+
 Value ofrToIndexValue(const OpFoldResult ofr, const Location loc,
                       OpBuilder &b) {
   if (Value val = dyn_cast<Value>(ofr)) {
@@ -344,17 +354,8 @@ OpFoldResult selectOFRs(const OpFoldResult condOFR, const OpFoldResult trueOFR,
                         OpBuilder &b) {
   auto trueValue = ofrToIndexValue(trueOFR, loc, b);
   auto falseValue = ofrToIndexValue(falseOFR, loc, b);
-  auto condValue = ofrToIndexValue(condOFR, loc, b);
-
-  // Ideally we should not be passing around everything as index type since mask
-  // analysis can come across i1 values, but that improvement is being left for
-  // future work. For now we just unwrap an index back into it's i1 value if
-  // necessary.
-  if (!condValue.getType().isInteger(1)) {
-    assert(condValue.getDefiningOp<arith::IndexCastOp>());
-    condValue = condValue.getDefiningOp<arith::IndexCastOp>().getOperand();
-    assert(condValue.getType().isInteger(1));
-  }
+  auto condValue = ofrToValue(condOFR, loc, b);
+  assert(condValue.getType().isInteger(1) && "Condition for selectOp must be a bool type");
 
   auto selectOp =
       b.create<arith::SelectOp>(loc, condValue, trueValue, falseValue);