Skip to content

Commit eb8661a

Browse files
red1bluelostred1bluelost
committed
[MaskAnalysis] Fix implement of scalar conjunction in mask analysis.
1 parent 2e65dc9 commit eb8661a

5 files changed

Lines changed: 106 additions & 53 deletions

File tree

‎include/triton-shared/Analysis/OpFoldResultUtils.h‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,9 @@
88
#ifndef TRITON_ANALYSIS_OPFOLDRESULT_UTILS_H
99
#define TRITON_ANALYSIS_OPFOLDRESULT_UTILS_H
1010

11+
#include "mlir/Dialect/Arith/IR/Arith.h"
1112
#include "mlir/IR/Location.h"
1213
#include "mlir/IR/OpDefinition.h"
13-
#include "mlir/Dialect/Arith/IR/Arith.h"
1414

1515
#include <optional>
1616

@@ -55,17 +55,23 @@ OpFoldResult subOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
5555
// result is an Integer Attribtue. Otherwise, insert the arith.muli
5656
// instruction if needed and use its result Value.
5757
OpFoldResult mulOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
58-
const Location loc, OpBuilder &b);
58+
const Location loc, OpBuilder &b);
5959

6060
OpFoldResult minOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
6161
const Location loc, OpBuilder &b);
6262

6363
OpFoldResult maxOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
6464
const Location loc, OpBuilder &b);
6565

66+
OpFoldResult selectOFRs(const OpFoldResult cond, const OpFoldResult trueOFR,
67+
const OpFoldResult falseOFR, const Location loc,
68+
OpBuilder &b);
69+
6670
OpFoldResult compareOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
67-
const arith::CmpIPredicate pred, const OpFoldResult trueVal,
68-
const OpFoldResult falseVal, const Location loc, OpBuilder &b);
71+
const arith::CmpIPredicate pred,
72+
const OpFoldResult trueVal,
73+
const OpFoldResult falseVal, const Location loc,
74+
OpBuilder &b);
6975
} // namespace mlir
7076

7177
#endif

‎lib/Analysis/MaskAnalysis.cpp‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -220,25 +220,25 @@ LogicalResult MaskState::minStateScalar(const MaskState &lhsState,
220220
const MaskState &rhsState, Location loc,
221221
OpBuilder &builder) {
222222
if (lhsState.scalar && rhsState.scalar) {
223-
dims.push_back(minOFRs(lhsState.dims[0], rhsState.dims[0], loc, builder));
224-
} else if (lhsState.scalar) {
225-
for (uint32_t i = 0; i < rhsState.getRank(); i++) {
226-
auto lhsDim = lhsState.dims[0];
227-
auto rhsDim = rhsState.dims[i];
228-
dims.push_back(minOFRs(lhsDim, rhsDim, loc, builder));
229-
}
230-
} else if (rhsState.scalar) {
231-
for (uint32_t i = 0; i < lhsState.getRank(); i++) {
232-
auto lhsDim = lhsState.dims[i];
233-
auto rhsDim = rhsState.dims[0];
234-
dims.push_back(minOFRs(lhsDim, rhsDim, loc, builder));
235-
}
236-
} else {
223+
InFlightDiagnostic diag =
224+
emitError(loc) << "Unexpected case where both lhs and rhs are scalars";
225+
return failure();
226+
}
227+
if (!lhsState.scalar && !rhsState.scalar) {
237228
InFlightDiagnostic diag =
238229
emitError(loc)
239230
<< "Unexpected case where both lhs and rhs are not scalars";
240231
return failure();
241232
}
233+
234+
auto &scalarState = lhsState.scalar ? lhsState : rhsState;
235+
auto &nonScalarState = lhsState.scalar ? rhsState : lhsState;
236+
for (uint32_t i = 0; i < nonScalarState.getRank(); i++) {
237+
auto nonScalarDim = nonScalarState.dims[i];
238+
dims.push_back(selectOFRs(scalarState.scalar, nonScalarDim,
239+
builder.getZeroAttr(builder.getIndexType()), loc,
240+
builder));
241+
}
242242
return success();
243243
}
244244

‎lib/Analysis/OpFoldResultUtils.cpp‎

Lines changed: 50 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,8 @@ OpFoldResult expandOFRIndex(OpFoldResult ofr, OpFoldResult targetForTy,
104104

105105
Value v = dyn_cast<Value>(ofr);
106106
if (!v)
107-
v = b.create<arith::ConstantOp>(loc, cast<IntegerAttr>(cast<Attribute>(ofr)));
107+
v = b.create<arith::ConstantOp>(loc,
108+
cast<IntegerAttr>(cast<Attribute>(ofr)));
108109

109110
Type ty = v.getType();
110111
if (targetTy == ty)
@@ -126,7 +127,8 @@ OpFoldResult expandOFRIndex(OpFoldResult ofr, OpFoldResult targetForTy,
126127
// This path is for case like:
127128
// input_ptr + (row_indices[:, None] + row_offsets[:,None] % mod_offset) *
128129
// stride_m + col_offsets[None, :] * stride_n
129-
// The modulo will be in shape of [ROW_SIZE, 1] while row_indices is in shape of [ROW_SIZE,].
130+
// The modulo will be in shape of [ROW_SIZE, 1] while row_indices is in
131+
// shape of [ROW_SIZE,].
130132
LLVM_DEBUG({
131133
llvm::dbgs() << "Reshaping ";
132134
shapedTy.dump();
@@ -135,14 +137,15 @@ OpFoldResult expandOFRIndex(OpFoldResult ofr, OpFoldResult targetForTy,
135137
});
136138
SmallVector<Value> shapeValues;
137139
for (auto dim : targetShapedTy.getShape()) {
138-
shapeValues.push_back(b.create<arith::ConstantOp>(
139-
loc, b.getIndexAttr(dim)));
140+
shapeValues.push_back(
141+
b.create<arith::ConstantOp>(loc, b.getIndexAttr(dim)));
140142
}
141143
RankedTensorType targetShapeTensorTy = RankedTensorType::get(
142144
targetShapedTy.getShape().size(), b.getIndexType());
143145
auto shapeTensor = b.create<tensor::FromElementsOp>(
144146
loc, targetShapeTensorTy, shapeValues);
145-
return b.create<triton::ReshapeOp>(loc, targetTy, v, shapeTensor).getResult();
147+
return b.create<triton::ReshapeOp>(loc, targetTy, v, shapeTensor)
148+
.getResult();
146149
}
147150
if (isa<IndexType>(targetEltTy) || isa<IndexType>(eltTy)) {
148151
assert((isa<IntegerType>(targetEltTy) || isa<IntegerType>(eltTy)) &&
@@ -228,7 +231,7 @@ OpFoldResult subOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
228231
}
229232

230233
OpFoldResult mulOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
231-
const Location loc, OpBuilder &b) {
234+
const Location loc, OpBuilder &b) {
232235
auto lhsIntAttr = getIntAttr(lhs);
233236
auto rhsIntAttr = getIntAttr(rhs);
234237

@@ -336,44 +339,61 @@ OpFoldResult maxOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
336339
return maxOp.getResult();
337340
}
338341

342+
OpFoldResult selectOFRs(const OpFoldResult condOFR, const OpFoldResult trueOFR,
343+
const OpFoldResult falseOFR, const Location loc,
344+
OpBuilder &b) {
345+
auto trueValue = ofrToIndexValue(trueOFR, loc, b);
346+
auto falseValue = ofrToIndexValue(falseOFR, loc, b);
347+
auto condValue = ofrToIndexValue(condOFR, loc, b);
348+
349+
if (!condValue.getType().isInteger(1)) {
350+
assert(condValue.getDefiningOp<arith::IndexCastOp>());
351+
condValue = condValue.getDefiningOp<arith::IndexCastOp>().getOperand();
352+
assert(condValue.getType().isInteger(1));
353+
}
354+
355+
auto selectOp =
356+
b.create<arith::SelectOp>(loc, condValue, trueValue, falseValue);
357+
return selectOp.getResult();
358+
}
359+
339360
OpFoldResult compareOFRs(const OpFoldResult lhs, const OpFoldResult rhs,
340-
const arith::CmpIPredicate pred, const OpFoldResult trueOFR,
341-
const OpFoldResult falseOFR, const Location loc, OpBuilder &b) {
361+
const arith::CmpIPredicate pred,
362+
const OpFoldResult trueOFR,
363+
const OpFoldResult falseOFR, const Location loc,
364+
OpBuilder &b) {
342365
auto lhsIntAttr = getIntAttr(lhs);
343366
auto rhsIntAttr = getIntAttr(rhs);
344367

345368
// both lhs and rhs are constants, return the result directly
346369
if (lhsIntAttr && rhsIntAttr) {
347370
switch (pred) {
348-
case arith::CmpIPredicate::eq:
349-
return *lhsIntAttr == *rhsIntAttr ? trueOFR : falseOFR;
350-
case arith::CmpIPredicate::ne:
351-
return *lhsIntAttr != *rhsIntAttr ? trueOFR : falseOFR;
352-
case arith::CmpIPredicate::slt:
353-
case arith::CmpIPredicate::ult:
354-
return *lhsIntAttr < *rhsIntAttr ? trueOFR : falseOFR;
355-
case arith::CmpIPredicate::sle:
356-
case arith::CmpIPredicate::ule:
357-
return *lhsIntAttr <= *rhsIntAttr ? trueOFR : falseOFR;
358-
case arith::CmpIPredicate::sgt:
359-
case arith::CmpIPredicate::ugt:
360-
return *lhsIntAttr > *rhsIntAttr ? trueOFR : falseOFR;
361-
case arith::CmpIPredicate::sge:
362-
case arith::CmpIPredicate::uge:
363-
return *lhsIntAttr >= *rhsIntAttr ? trueOFR : falseOFR;
364-
default:
365-
llvm_unreachable("Unsupported predicate");
371+
case arith::CmpIPredicate::eq:
372+
return *lhsIntAttr == *rhsIntAttr ? trueOFR : falseOFR;
373+
case arith::CmpIPredicate::ne:
374+
return *lhsIntAttr != *rhsIntAttr ? trueOFR : falseOFR;
375+
case arith::CmpIPredicate::slt:
376+
case arith::CmpIPredicate::ult:
377+
return *lhsIntAttr < *rhsIntAttr ? trueOFR : falseOFR;
378+
case arith::CmpIPredicate::sle:
379+
case arith::CmpIPredicate::ule:
380+
return *lhsIntAttr <= *rhsIntAttr ? trueOFR : falseOFR;
381+
case arith::CmpIPredicate::sgt:
382+
case arith::CmpIPredicate::ugt:
383+
return *lhsIntAttr > *rhsIntAttr ? trueOFR : falseOFR;
384+
case arith::CmpIPredicate::sge:
385+
case arith::CmpIPredicate::uge:
386+
return *lhsIntAttr >= *rhsIntAttr ? trueOFR : falseOFR;
387+
default:
388+
llvm_unreachable("Unsupported predicate");
366389
}
367390
}
368391

369392
auto lhsValue = ofrToIndexValue(lhs, loc, b);
370393
auto rhsValue = ofrToIndexValue(rhs, loc, b);
371-
auto trueValue = ofrToIndexValue(trueOFR, loc, b);
372-
auto falseValue = ofrToIndexValue(falseOFR, loc, b);
373394

374395
auto cmpOp = b.create<arith::CmpIOp>(loc, pred, lhsValue, rhsValue);
375-
auto selectOp = b.create<arith::SelectOp>(loc, cmpOp, trueValue, falseValue);
376-
return selectOp.getResult();
396+
return selectOFRs(cmpOp.getResult(), trueOFR, falseOFR, loc, b);
377397
}
378398

379399
} // namespace mlir

‎python/examples/test_mask.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,3 +37,29 @@ def test(in0, out0):
3737
print(input)
3838
print(output)
3939
torch.testing.assert_close(output, torch.tensor([-1, -1, -1, -1, -2, -2, -2, -2], device=device, dtype=torch.int32))
40+
41+
42+
def test_mask_with_scalar_in_conjunction(device):
43+
if device == 'cpu':
44+
triton.runtime.driver.set_active(CPUDriver())
45+
46+
@triton.jit
47+
def kernel(in0, out0, mask, value):
48+
offs = tl.arange(0, 8)
49+
out_offs = tl.arange(0, 8)
50+
a = tl.load(in0 + offs, mask=(value < 5) & (offs < mask), other=-1)
51+
tl.store(out0 + out_offs, a)
52+
53+
# Test scalar mask evaluate to True
54+
SIZE = 8
55+
input = torch.arange(0, SIZE, device=device, dtype=torch.int32)
56+
output = torch.full((SIZE,), -2, device=device, dtype=torch.int32)
57+
kernel[(1,)](input, output, 4, 3)
58+
torch.testing.assert_close(output, torch.tensor([0, 1, 2, 3, -1, -1, -1, -1], device=device, dtype=torch.int32))
59+
60+
# Test scalar mask evaluate to False
61+
SIZE = 8
62+
input = torch.arange(0, SIZE, device=device, dtype=torch.int32)
63+
output = torch.full((SIZE,), -2, device=device, dtype=torch.int32)
64+
kernel[(1,)](input, output, 4, 8)
65+
torch.testing.assert_close(output, torch.tensor([-1, -1, -1, -1, -1, -1, -1, -1], device=device, dtype=torch.int32))

‎test/Conversion/TritonToStructured/mask_ld_st_scalar_dim.mlir‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// RUN: triton-shared-opt --triton-to-structured --remove-dead-values --canonicalize %s | FileCheck %s
1+
// RUN: triton-shared-opt --triton-to-structured --remove-dead-values --canonicalize --cse %s | FileCheck %s
22

33
module {
44
tt.func @mask_ld_st_scalar(
@@ -38,10 +38,11 @@ module {
3838
}
3939
}
4040

41-
// CHECK: %{{.*}} = "tts.load"(%{{.*}}, %{{.*}}) <{operandSegmentSizes = array<i32: 1, 1, 0>, static_mask_dims = array<i64: -9223372036854775808, 1>}> : (tensor<2x1x!tt.ptr<f32>>, index) -> tensor<2x1xf32>
42-
// CHECK: %{{.*}} = "tts.load"(%{{.*}}, %{{.*}}) <{operandSegmentSizes = array<i32: 1, 1, 0>, static_mask_dims = array<i64: -9223372036854775808, 1>}> : (tensor<2x1x!tt.ptr<f32>>, index) -> tensor<2x1xf32>
43-
// CHECK: "tts.store"(%{{.*}}, %{{.*}}, %{{.*}}) <{static_mask_dims = array<i64: -9223372036854775808, 1>}> : (tensor<2x1x!tt.ptr<f32>>, tensor<2x1xf32>, index) -> ()
44-
// CHECK: "tts.store"(%{{.*}}, %{{.*}}, %{{.*}}) <{static_mask_dims = array<i64: -9223372036854775808, 1>}> : (tensor<2x1x!tt.ptr<f32>>, tensor<2x1xf32>, index) -> ()
41+
// CHECK: %{{.*}} = "tts.load"(%{{.*}}, %{{.*}}, %{{.*}}) <{operandSegmentSizes = array<i32: 1, 2, 0>, static_mask_dims = array<i64: -9223372036854775808, -9223372036854775808>}> : (tensor<2x1x!tt.ptr<f32>>, index, index) -> tensor<2x1xf32>
42+
// CHECK: %{{.*}} = "tts.load"(%{{.*}}, %{{.*}}, %{{.*}}) <{operandSegmentSizes = array<i32: 1, 2, 0>, static_mask_dims = array<i64: -9223372036854775808, -9223372036854775808>}> : (tensor<2x1x!tt.ptr<f32>>, index, index) -> tensor<2x1xf32>
43+
// CHECK: "tts.store"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) <{static_mask_dims = array<i64: -9223372036854775808, -9223372036854775808>}> : (tensor<2x1x!tt.ptr<f32>>, tensor<2x1xf32>, index, index) -> ()
44+
// CHECK: "tts.store"(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) <{static_mask_dims = array<i64: -9223372036854775808, -9223372036854775808>}> : (tensor<2x1x!tt.ptr<f32>>, tensor<2x1xf32>, index, index) -> ()
45+
4546

4647
// Original Triton Function:
4748
// def test_masked_ld_st(

0 commit comments

Comments
 (0)