Skip to content

Siwasaki/pr/loadstorescalar #16

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
shintaro-iwasaki wants to merge 7 commits into
base: triton-mlir
Choose a base branch
from
Open

Siwasaki/pr/loadstorescalar #16

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
5898352
[Triton-IR] Fix LoadOp definition (#771) (#777)
shintaro-iwasaki Oct 14, 2022
e948a61
[Triton-MLIR] fix a tiny bug in coalesce pass (#782)
goostavz Oct 17, 2022
38a8066
[OPTIMIZER] Updated TritonGPU-combine pass (#784)
ptillet Oct 17, 2022
4464646
[Triton-MLIR][BACKEND] Fix masked load store op vector size (#785)
Superjomn Oct 18, 2022
6f1611f
[TritonIR] Update Load/StoreOps to support scalar values
shintaro-iwasaki Oct 17, 2022
db1c304
[TritonIR] tl.reduce returns a scalar value if input is a 1D tensor
shintaro-iwasaki Oct 17, 2022
07c440e
[Tests] Add tests to check type inference
shintaro-iwasaki Oct 17, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 6 additions & 7 deletions include/triton/Dialect/Triton/IR/TritonOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,11 +105,10 @@ def TT_AddPtrOp : TT_Op<"addptr",
def TT_LoadOp : TT_Op<"load",
[SameOperandsAndResultShape,
SameOperandsAndResultEncoding,
SameVariadicOperandSize,
AttrSizedOperandSegments,
MemoryEffects<[MemRead]>,
TypesMatchWith<"infer ptr type from result type",
"result", "ptr",
"getPointerTypeFromTensor($_self)">,
"result", "ptr", "getPointerTypeSameShape($_self)">,
TypesMatchWith<"infer mask type from result type or none",
"result", "mask", "getI1SameShape($_self)",
"($_op.getOperands().size() <= 1) || std::equal_to<>()">,
Expand All @@ -118,7 +117,7 @@ def TT_LoadOp : TT_Op<"load",
"($_op.getOperands().size() <= 2) || std::equal_to<>()">]> {
let summary = "load";

let arguments = (ins TT_PtrTensor:$ptr, Optional<I1Tensor>:$mask, Optional<TT_Type>:$other,
let arguments = (ins TT_PtrLike:$ptr, Optional<TT_BoolLike>:$mask, Optional<TT_Type>:$other,
TT_CacheModifierAttr:$cache, TT_EvictionPolicyAttr:$evict,
BoolAttr:$isVolatile);

Expand Down Expand Up @@ -147,13 +146,13 @@ def TT_StoreOp : TT_Op<"store",
MemoryEffects<[MemWrite]>,
TypesMatchWith<"infer ptr type from value type",
"value", "ptr",
"getPointerTypeFromTensor($_self)">,
"getPointerTypeSameShape($_self)">,
TypesMatchWith<"infer mask type from value type",
"value", "mask", "getI1SameShape($_self)",
"($_op.getOperands().size() <= 2) || std::equal_to<>()">]> {
let summary = "store";

let arguments = (ins TT_PtrTensor:$ptr, TT_Type:$value, Optional<I1Tensor>:$mask);
let arguments = (ins TT_PtrLike:$ptr, TT_Type:$value, Optional<TT_BoolLike>:$mask);

let builders = [
OpBuilder<(ins "Value":$ptr, "Value":$value)>,
Expand Down Expand Up @@ -318,7 +317,7 @@ def TT_ReduceOp : TT_Op<"reduce", [NoSideEffect,

let arguments = (ins TT_RedOpAttr:$redOp, TT_Tensor:$operand, I32Attr:$axis);

let results = (outs TT_Tensor:$result);
let results = (outs TT_Type:$result);

let builders = [
OpBuilder<(ins "triton::RedOp":$redOp, "Value":$operand, "int":$axis)>,
Expand Down
30 changes: 29 additions & 1 deletion lib/Analysis/AxisInfo.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,33 @@ ChangeResult AxisInfoAnalysis::visitOperation(
}
curr = AxisInfo(contiguity, divisibility, constancy);
}

// CmpI
if ((llvm::dyn_cast<arith::CmpIOp>(op) ||
llvm::dyn_cast<triton::gpu::CmpIOp>(op)) &&
op->getResult(0).getType().dyn_cast<TensorType>()) {
auto resTy = op->getResult(0).getType().cast<TensorType>();
short rank = resTy.getRank();
auto lhsInfo = operands[0]->getValue();
auto rhsInfo = operands[1]->getValue();
auto shape = resTy.getShape();

AxisInfo::DimVectorT contiguity, divisibility, constancy;
for (short d = 0; d < rank; ++d) {
if (rhsInfo.getConstancy(d) % lhsInfo.getContiguity(d) == 0 ||
rhsInfo.getConstancy(d) % lhsInfo.getConstancy(d))
constancy.push_back(
gcd(lhsInfo.getDivisibility(d), rhsInfo.getDivisibility(d)));
else
constancy.push_back(1);

divisibility.push_back(shape[d]);
contiguity.push_back(1);
}

curr = AxisInfo(contiguity, divisibility, constancy);
}

// UnrealizedConversionCast
// This is needed by TritonGPUToLLVM, to get AxisInfo when the graph is
// in the process of a PartialConversion, where UnrealizedConversionCast
Expand All @@ -219,7 +246,8 @@ ChangeResult AxisInfoAnalysis::visitOperation(
if (curr.getRank() == 0) {
return markAllPessimisticFixpoint(op->getResults());
}
// join all latice elements

// join all lattice elements
ChangeResult result = ChangeResult::NoChange;
for (Value value : op->getResults()) {
result |= getLatticeElement(value).join(curr);
Expand Down
Loading