Skip to content

Commit 3596dc5

Browse files
jax-triton-devreichlfl
jax-triton-dev
authored and
reichlfl
committed
OpenXLA-specific changes
1 parent 174e780 commit 3596dc5

File tree

28 files changed

+2109
-15
lines changed

28 files changed

+2109
-15
lines changed

‎BUILD

+905
Large diffs are not rendered by default.

‎lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp

+2-1
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,8 @@ SmallVector<Value> reorderValues(const SmallVector<Value> &values, Type inType,
4040
auto ouEltTy = ouTensorTy.getElementType();
4141
if (inBitWidth == ouBitWidth)
4242
return values;
43-
if (inBitWidth == 16 && ouBitWidth == 32) {
43+
if ((inBitWidth == 16 && ouBitWidth == 32) ||
44+
(inBitWidth == 32 && ouBitWidth == 16)) {
4445
SmallVector<Value> ret;
4546
for (unsigned i = 0; i < values.size(); i += 8) {
4647
ret.push_back(values[i]);

‎lib/Conversion/TritonGPUToLLVM/ViewOpToLLVM.cpp

+2-1
Original file line numberDiff line numberDiff line change
@@ -87,8 +87,9 @@ struct ArithConstantSplatOpConversion
8787
// LLVM IR.
8888
if (type::isFloat8(elemType))
8989
elemType = rewriter.getIntegerType(8);
90-
auto constOp = rewriter.create<LLVM::ConstantOp>(loc, elemType, val);
9190
auto typeConverter = getTypeConverter();
91+
auto constOp = rewriter.create<LLVM::ConstantOp>(
92+
loc, typeConverter->convertType(elemType), val);
9293
auto llStruct = SplatOpConversion::convertSplatLikeOp(
9394
elemType, op.getType(), constOp, typeConverter, rewriter, loc);
9495
rewriter.replaceOp(op, llStruct);

‎lib/Dialect/TritonGPU/IR/Dialect.cpp

+5
Original file line numberDiff line numberDiff line change
@@ -2717,6 +2717,11 @@ struct CanonicalizeConvertFromAlloc
27172717
auto convert = op.getSrc().getDefiningOp<ConvertLayoutOp>();
27182718
if (!convert)
27192719
return failure();
2720+
// LocalAllocOp lowering doesn't support going from DotOperandEncoding
2721+
// to SharedEncoding, so we want to keep this layout conversion.
2722+
if (mlir::isa<triton::gpu::DotOperandEncodingAttr>(
2723+
convert.getSrc().getType().getEncoding()))
2724+
return failure();
27202725
rewriter.replaceOpWithNewOp<triton::gpu::LocalAllocOp>(
27212726
op, op->getResult(0).getType(), convert.getSrc());
27222727
return mlir::success();

‎lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp

+24
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,21 @@ static Value getSharedMemoryMMAOperand(Value v, mlir::PatternRewriter &rewriter,
153153
auto newType = MemDescType::get(argType.getShape(), argType.getElementType(),
154154
newLayout, SharedMemorySpace);
155155
rewriter.setInsertionPointAfterValue(arg);
156+
157+
// LocalAllocOp lowering doesn't support going from DotOperandEncoding
158+
// to SharedEncoding.
159+
if (auto dotOpEnc = mlir::dyn_cast<DotOperandEncodingAttr>(
160+
argType.getEncoding())) {
161+
// Create a layout conversion from DotOperandEncoding to BlockedEncoding
162+
// then pass it to the LocalAllocOp.
163+
auto newArgType = RankedTensorType::get(
164+
argType.getShape(), argType.getElementType(), dotOpEnc.getParent());
165+
auto dotOperandToBlockedCvt =
166+
rewriter.create<ConvertLayoutOp>(arg.getLoc(), newArgType, arg);
167+
return rewriter.create<LocalAllocOp>(arg.getLoc(), newType,
168+
dotOperandToBlockedCvt);
169+
}
170+
156171
return rewriter.create<LocalAllocOp>(arg.getLoc(), newType, arg);
157172
}
158173

@@ -162,6 +177,15 @@ class BlockedToMMA : public mlir::OpRewritePattern<DotOp> {
162177
mutable llvm::DenseMap<Operation *, unsigned> dotOpInstNs;
163178

164179
static bool bwdFilter(Operation *op) {
180+
// Dot operand layout assignment to Predicates are not currently supported
181+
// during lowering from TritonGPU to LLVM in Triton for MMA cases. This
182+
// condition limits visibility of the original bit-width so that predicate
183+
// are not considered, hence, kwidth can never be = 32.
184+
if (isa<arith::UIToFPOp>(op)) {
185+
Type srcType = getElementTypeOrSelf(op->getOperand(0));
186+
if (srcType.isInteger(1))
187+
return false;
188+
}
165189
return op->getNumOperands() == 1 &&
166190
(isa<FpToFpOp, BitcastOp, ConvertLayoutOp>(op) ||
167191
isPureUnaryInlineAsm(op) ||

‎lib/Dialect/TritonGPU/Transforms/OptimizeDotOperands.cpp

+16-1
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,8 @@ class HoistLayoutConversion : public OpRewritePattern<ConvertLayoutOp> {
111111
PatternRewriter &rewriter) const override {
112112
// Only consider conversions to dot operand.
113113
auto cvtTy = cast<RankedTensorType>(cvt.getType());
114-
if (!isa<DotOperandEncodingAttr>(cvtTy.getEncoding()))
114+
auto dotOpEnc = dyn_cast<DotOperandEncodingAttr>(cvtTy.getEncoding());
115+
if (!dotOpEnc)
115116
return failure();
116117

117118
auto src = cvt.getSrc().getDefiningOp();
@@ -126,6 +127,12 @@ class HoistLayoutConversion : public OpRewritePattern<ConvertLayoutOp> {
126127
[](Type ty) { return isa<RankedTensorType>(ty); }))
127128
return failure();
128129

130+
// Quick handling to fix loading issues when computing the original
131+
// bitwidth is unable to realize that there is a mixed-precision dot
132+
// (hence kWidth = 1) but wants to hoist through the type conversion.
133+
if (isa<arith::ExtFOp>(src) && dotOpEnc.getKWidth() == 1)
134+
return failure();
135+
129136
// Only consider custom conversions or arith ops.
130137
// TODO(jlebar): Is this too restrictive?
131138
if (!isa<FpToFpOp, BitcastOp>(src) && !isPureUnaryInlineAsm(src) &&
@@ -138,6 +145,14 @@ class HoistLayoutConversion : public OpRewritePattern<ConvertLayoutOp> {
138145
if (isa<arith::TruncIOp, arith::TruncFOp, arith::SelectOp>(src))
139146
return failure();
140147

148+
// Don't hoist through u1 -> fp casts as they aren't supported in
149+
// ElementwiseOpToLLVM::reorderValues().
150+
if (isa<arith::UIToFPOp>(src)) {
151+
Type srcType = getElementTypeOrSelf(src->getOperand(0));
152+
if (srcType.isInteger(1))
153+
return failure();
154+
}
155+
141156
// Check that the conversion is transitively dependent on a load, and all
142157
// operations between the load and the conversion are layout preserving.
143158
//

‎lib/Dialect/TritonGPU/Transforms/Prefetch.cpp

+16-1
Original file line numberDiff line numberDiff line change
@@ -140,8 +140,14 @@ Value Prefetcher::generatePrefetch(Value v, unsigned opIdx, bool isPrologue,
140140
type.getMemorySpace()),
141141
v, offsetsVal);
142142

143+
// We need to assign kwidth to zero in the case where the parent layout is
144+
// Blocked, otherwise the verifier emits a failure. The parent layout is
145+
// Blocked only when Tensor Cores are disabled.
146+
int kwidth = dyn_cast<triton::gpu::BlockedEncodingAttr>(dotEncoding)
147+
? 0
148+
: prefetchWidth / 8;
143149
auto dotOperandEnc = triton::gpu::DotOperandEncodingAttr::get(
144-
builder.getContext(), opIdx, dotEncoding, prefetchWidth / 8);
150+
builder.getContext(), opIdx, dotEncoding, kwidth);
145151
Value prefetchSlice = builder.create<triton::gpu::LocalLoadOp>(
146152
v.getLoc(), RankedTensorType::get(shape, elementType, dotOperandEnc),
147153
newSmem);
@@ -187,6 +193,15 @@ LogicalResult Prefetcher::initialize() {
187193
break;
188194
if (!op->getResult(0).hasOneUse())
189195
break;
196+
// Similar to issues faced in HoistLayoutConversion pattern in
197+
// OptimizeDotOperands.cpp, we can't propagate through type casts from
198+
// predicates as they aren't supported in Triton when encoded with dot_op
199+
// layout.
200+
if (isa<arith::UIToFPOp>(op)) {
201+
Type srcType = getElementTypeOrSelf(op->getOperand(0));
202+
if (srcType.isInteger(1))
203+
break;
204+
}
190205
rets.push_back(op->getOperand(0));
191206
if (auto cvt = dyn_cast<triton::gpu::LocalLoadOp>(op)) {
192207
foundConvertFromShared = true;

‎python/BUILD

+77
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
# NOTE: Do not depend on any targets from this directory,
2+
# but use //third_party/py/triton instead.
3+
4+
load("@pybind11_bazel//:build_defs.bzl", "pybind_extension")
5+
6+
package(
7+
default_applicable_licenses = ["//:license"],
8+
default_visibility = [
9+
"//third_party/py/triton:__pkg__",
10+
"@triton//python:__subpackages__",
11+
],
12+
)
13+
14+
cc_library(
15+
name = "passes",
16+
hdrs = ["src/passes.h"],
17+
includes = ["src"],
18+
visibility = ["@triton//third_party:__subpackages__"],
19+
)
20+
21+
pybind_extension(
22+
name = "libtriton",
23+
srcs = [
24+
"src/interpreter.cc",
25+
"src/ir.cc",
26+
"src/llvm.cc",
27+
"src/main.cc",
28+
"src/passes.cc",
29+
],
30+
copts = ["-DTRITON_BACKENDS_TUPLE=(nvidia)"],
31+
deps = [
32+
":passes",
33+
"@llvm-project//llvm:Core",
34+
"@llvm-project//llvm:IPO",
35+
"@llvm-project//llvm:IRReader",
36+
"@llvm-project//llvm:InstCombine",
37+
"@llvm-project//llvm:Linker",
38+
"@llvm-project//llvm:MC",
39+
"@llvm-project//llvm:Passes",
40+
"@llvm-project//llvm:Support",
41+
"@llvm-project//llvm:Target",
42+
"@llvm-project//mlir:BuiltinToLLVMIRTranslation",
43+
"@llvm-project//mlir:BytecodeWriter",
44+
"@llvm-project//mlir:ControlFlowDialect",
45+
"@llvm-project//mlir:ConversionPasses",
46+
"@llvm-project//mlir:IR",
47+
"@llvm-project//mlir:IndexDialect",
48+
"@llvm-project//mlir:LLVMDialect",
49+
"@llvm-project//mlir:LLVMIRTransforms",
50+
"@llvm-project//mlir:LLVMToLLVMIRTranslation",
51+
"@llvm-project//mlir:NVVMToLLVMIRTranslation",
52+
"@llvm-project//mlir:Parser",
53+
"@llvm-project//mlir:Pass",
54+
"@llvm-project//mlir:Support",
55+
"@llvm-project//mlir:ToLLVMIRTranslation",
56+
"@llvm-project//mlir:Transforms",
57+
"//:TritonAnalysis",
58+
"//:TritonDialects",
59+
"//:TritonGPUToLLVM",
60+
"//:TritonGPUTransforms",
61+
"//:TritonHSACO",
62+
"//:TritonLLVMIR",
63+
"//:TritonNvidiaGPUTransforms",
64+
"//:TritonPTX",
65+
"//:TritonToTritonGPU",
66+
"//:TritonTools",
67+
"//:TritonTransforms",
68+
"@triton//third_party/nvidia:triton_nvidia",
69+
],
70+
)
71+
72+
filegroup(
73+
name = "files",
74+
srcs = glob(
75+
include = ["triton/**/*.py"],
76+
),
77+
)

‎python/src/ir.cc

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
#include <pybind11/functional.h>
1+
#include <pybind11/functional.h>
22
#include <pybind11/pybind11.h>
33
#include <pybind11/stl.h>
44

‎python/test/regression/BUILD

+26
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
load("//third_party/py/pytest:pytest_defs.bzl", "pytest_multi_tests")
2+
3+
package(
4+
default_applicable_licenses = ["//:license"],
5+
)
6+
7+
pytest_multi_tests(
8+
name = "tests",
9+
size = "large",
10+
srcs = ["conftest.py"],
11+
shard_count = 10,
12+
tags = [
13+
"config-cuda-only",
14+
"requires-gpu-sm80",
15+
],
16+
tests = glob(
17+
include = ["test_*.py"],
18+
exclude = [
19+
"test_performance.py", #TODO(b/321005767): fix failing test
20+
],
21+
),
22+
deps = [
23+
"//third_party/py/torch:pytorch",
24+
"//third_party/py/triton",
25+
],
26+
)

‎python/test/regression/conftest.py

+12
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
# content of conftest.py
2+
3+
import pytest
4+
5+
6+
def pytest_addoption(parser):
7+
parser.addoption("--device", action="store", default='cuda')
8+
9+
10+
@pytest.fixture
11+
def device(request):
12+
return request.config.getoption("--device")

0 commit comments

Comments
 (0)