Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
2 changes: 1 addition & 1 deletion cmake/llvm-hash.txt
Original file line number Diff line number Diff line change
@@ -1 +1 @@
1f20eee6dc367bd202895e3eedb03974a628ef16
86b69c31642e98f8357df62c09d118ad1da4e16a
4 changes: 0 additions & 4 deletions include/triton/Dialect/Triton/IR/TritonOps.td
Original file line number Diff line number Diff line change
Expand Up @@ -727,10 +727,6 @@ def TT_ReduceOp: TT_Op<"reduce",
llvm::SmallVector<RankedTensorType> getInputTypes();
llvm::SmallVector<Type> getElementTypes();
unsigned getNumOperands();

// Returns the CombineOp iff this ReduceOp's region contains only
// one CombineOp other than the return, or nullptr if not applicable.
::mlir::Operation *getSingleCombiner();
}];
}

Expand Down
16 changes: 0 additions & 16 deletions lib/Dialect/Triton/IR/Ops.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -503,22 +503,6 @@ llvm::SmallVector<Type> ReduceOp::getElementTypes() {
return getElementTypesImpl(this->getOperands());
}

::mlir::Operation *ReduceOp::getSingleCombiner() {
if (getNumOperands() != 1 || getNumResults() != 1)
return nullptr;
Block *block = &(*getCombineOp().begin());
Operation *yield = block->getTerminator();
Operation *reduceOp = yield->getOperand(0).getDefiningOp();
if (!reduceOp || reduceOp->getNumOperands() != 2 ||
reduceOp->getNumResults() != 1)
return nullptr;
if (reduceOp->getOperand(0) != block->getArgument(0) ||
reduceOp->getOperand(1) != block->getArgument(1))
return nullptr;

return reduceOp;
}

unsigned ReduceOp::getNumOperands() { return this->getOperands().size(); }

//-- ScanOp --
Expand Down
4 changes: 2 additions & 2 deletions lib/Dialect/Triton/Transforms/Combine.td
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ def CombineDotAddIPattern : Pat<
[(Constraint<CPred<"isZero($0)">> $c),
(Constraint<CPred<"res->hasOneUse()">, "dot result has a single use">)]>;
def CombineDotAddFPattern : Pat<
(Arith_AddFOp $d, (TT_DotOp:$res $a, $b, $c, $inputPrecision, $maxNumImpreciseAcc), $fastmath, $denorm),
(Arith_AddFOp $d, (TT_DotOp:$res $a, $b, $c, $inputPrecision, $maxNumImpreciseAcc), $fastmath),
(TT_DotOp $a, $b, $d, $inputPrecision, $maxNumImpreciseAcc, (location $res)),
[(Constraint<CPred<"isZero($0)">> $c),
(Constraint<CPred<"::llvm::cast<::mlir::IntegerAttr>($0).getInt() == 0">> $maxNumImpreciseAcc),
Expand All @@ -29,7 +29,7 @@ def CombineDotAddIRevPattern : Pat<
[(Constraint<CPred<"isZero($0)">> $c),
(Constraint<CPred<"res->hasOneUse()">, "dot result has a single use">)]>;
def CombineDotAddFRevPattern : Pat<
(Arith_AddFOp (TT_DotOp:$res $a, $b, $c, $inputPrecision, $maxNumImpreciseAcc), $d, $fastmath, $denorm),
(Arith_AddFOp (TT_DotOp:$res $a, $b, $c, $inputPrecision, $maxNumImpreciseAcc), $d, $fastmath),
(TT_DotOp $a, $b, $d, $inputPrecision, $maxNumImpreciseAcc, (location $res)),
[(Constraint<CPred<"isZero($0)">> $c),
(Constraint<CPred<"::llvm::cast<::mlir::IntegerAttr>($0).getInt() == 0">> $maxNumImpreciseAcc),
Expand Down
108 changes: 0 additions & 108 deletions test/Conversion/amd/tritongpu_to_llvm.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -62,111 +62,3 @@ module attributes {"triton_gpu.num-ctas" = 1 : i32, "triton_gpu.num-warps" = 4 :
tt.return
}
}

// -----

#blocked1 = #triton_gpu.blocked<{sizePerThread = [2], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>
module attributes {"triton_gpu.num-ctas" = 1 : i32, "triton_gpu.num-warps" = 4 : i32} {
// CHECK-LABEL: atomic_add_f16
tt.func @atomic_add_f16(%arg0: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg1 : tensor<256xi1, #blocked1>, %arg2 : tensor<256xf16, #blocked1>) {
%range = tt.make_range {end = 256 : i32, start = 0 : i32} : tensor<256xi32, #blocked1>
%base_ptr = tt.splat %arg0 : !tt.ptr<f16> -> tensor<256x!tt.ptr<f16>, #blocked1>
%ptr = tt.addptr %base_ptr, %range : tensor<256x!tt.ptr<f16>, #blocked1>, tensor<256xi32, #blocked1>
// CHECK: llvm.cond_br
// CHECK: llvm.atomicrmw fadd {{.*}} vector<2xf16>
%0 = tt.atomic_rmw fadd, relaxed, gpu, %ptr, %arg2, %arg1 : (tensor<256x!tt.ptr<f16>, #blocked1>, tensor<256xf16, #blocked1>, tensor<256xi1, #blocked1>) -> tensor<256xf16, #blocked1>
tt.return
}
}

// -----

#blocked2 = #triton_gpu.blocked<{sizePerThread = [2], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>
module attributes {"triton_gpu.num-ctas" = 1 : i32, "triton_gpu.num-warps" = 4 : i32} {
// CHECK-LABEL: atomic_add_bf16
tt.func @atomic_add_bf16(%arg0: !tt.ptr<bf16> {tt.divisibility = 16 : i32}, %arg1 : tensor<256xi1, #blocked2>, %arg2 : tensor<256xbf16, #blocked2>) {
%range = tt.make_range {end = 256 : i32, start = 0 : i32} : tensor<256xi32, #blocked2>
%base_ptr = tt.splat %arg0 : !tt.ptr<bf16> -> tensor<256x!tt.ptr<bf16>, #blocked2>
%ptr = tt.addptr %base_ptr, %range : tensor<256x!tt.ptr<bf16>, #blocked2>, tensor<256xi32, #blocked2>
// CHECK: llvm.cond_br
// CHECK: llvm.atomicrmw fadd {{.*}} vector<2xbf16>
%0 = tt.atomic_rmw fadd, relaxed, gpu, %ptr, %arg2, %arg1 : (tensor<256x!tt.ptr<bf16>, #blocked2>, tensor<256xbf16, #blocked2>, tensor<256xi1, #blocked2>) -> tensor<256xbf16, #blocked2>
tt.return
}
}

// -----

#blocked3 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [64], warpsPerCTA = [1], order = [0]}>
module attributes {"triton_gpu.num-ctas" = 1 : i32, "triton_gpu.num-warps" = 1 : i32, "triton_gpu.threads-per-warp" = 64 : i32} {
// CHECK-LABEL: reduce_dpp_max
tt.func @reduce_dpp_max(%arg0: tensor<64xf32, #blocked3>) {
// CHECK: rocdl.update.dpp
// CHECK-SAME: with 280, 15, 15, true : f32
// CHECK-NEXT: llvm.intr.maxnum

// CHECK-NEXT: rocdl.update.dpp
// CHECK-SAME: with 276, 15, 15, true : f32
// CHECK-NEXT: llvm.intr.maxnum

// CHECK-NEXT: rocdl.update.dpp
// CHECK-SAME: with 274, 15, 15, true : f32
// CHECK-NEXT: llvm.intr.maxnum

// CHECK-NEXT: rocdl.update.dpp
// CHECK-SAME: with 273, 15, 15, true : f32
// CHECK-NEXT: llvm.intr.maxnum

// CHECK-NEXT: rocdl.update.dpp
// CHECK-SAME: with 322, 10, 15, true : f32
// CHECK-NEXT: llvm.intr.maxnum

// CHECK-NEXT: rocdl.update.dpp
// CHECK-SAME: with 323, 15, 15, true : f32
// CHECK-NEXT: llvm.intr.maxnum

// CHECK: llvm.amdgcn.readlane
%0 = "tt.reduce"(%arg0) <{axis = 0 : i32}> ({
^bb0(%arg1: f32, %arg2: f32):
%1 = arith.maxnumf %arg1, %arg2 : f32
tt.reduce.return %1 : f32
}) : (tensor<64xf32, #blocked3>) -> f32
tt.return
}
}

// -----

#blocked4 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [64], warpsPerCTA = [1], order = [0]}>
module attributes {"triton_gpu.num-ctas" = 1 : i32, "triton_gpu.num-warps" = 1 : i32, "triton_gpu.threads-per-warp" = 64 : i32} {
// CHECK-LABEL: reduce_xor_max
tt.func @reduce_xor_max(%arg0: tensor<32xf32, #blocked4>) {
// CHECK: rocdl.ds_swizzle
// CHECK: llvm.intr.maxnum

// CHECK: rocdl.update.dpp
// CHECK-SAME: with 280, 15, 12, false : i32
// CHECK: rocdl.update.dpp
// CHECK-SAME: with 264, 15, 3, false : i32
// CHECK: llvm.intr.maxnum

// CHECK: rocdl.update.dpp
// CHECK-SAME: with 276, 15, 10, false : i32
// CHECK: rocdl.update.dpp
// CHECK-SAME: with 260, 15, 5, false : i32
// CHECK: llvm.intr.maxnum

// CHECK: rocdl.update.dpp
// CHECK-SAME: with 78, 15, 15, false : i32
// CHECK: llvm.intr.maxnum

// CHECK: rocdl.update.dpp
// CHECK-SAME: with 177, 15, 15, false : i32
%0 = "tt.reduce"(%arg0) <{axis = 0 : i32}> ({
^bb0(%arg1: f32, %arg2: f32):
%1 = arith.maxnumf %arg1, %arg2 : f32
tt.reduce.return %1 : f32
}) : (tensor<32xf32, #blocked4>) -> f32
tt.return
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@
#include "mlir/IR/Dialect.h"
#include "mlir/IR/PatternMatch.h"
#include "triton/Dialect/Triton/IR/Traits.h"

// clang-format off
#include "amd/include/Dialect/TritonAMDGPU/IR/Dialect.h.inc"
// clang-format on
Expand Down
11 changes: 0 additions & 11 deletions third_party/amd/include/TritonAMDGPUToLLVM/TargetUtils.h
Original file line number Diff line number Diff line change
Expand Up @@ -19,17 +19,6 @@ enum class ISAFamily {
// Deduces the corresponding ISA family for the given target gfx |arch|.
ISAFamily deduceISAFamily(llvm::StringRef arch);

// Here is a partial definition of DppCtrl enums. For the complete definition,
// please check:
// https://github.com/llvm/llvm-project/blob/8c75290/llvm/lib/Target/AMDGPU/SIDefines.h#L939
enum class DppCtrl : uint32_t {
QUAD_PERM_FIRST = 0,
ROW_SHL0 = 0x100,
ROW_SHR0 = 0x110,
BCAST15 = 0x142,
BCAST31 = 0x143
};

} // namespace mlir::triton::AMD

#endif // TRITON_CONVERSION_TRITONGPU_TO_LLVM_TARGETUTILS_H
41 changes: 21 additions & 20 deletions third_party/amd/lib/TritonAMDGPUToLLVM/LoadStoreOpToLLVM.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -768,11 +768,7 @@ struct AtomicRMWOpConversion
// tensor
if (tensorTy) {
auto valTy = cast<RankedTensorType>(val.getType());
Type elTy = valTy.getElementType();
vec = std::min<unsigned>(vec, llvm::isa<FloatType>(elTy) &&
elTy.getIntOrFloatBitWidth() == 16
? 2
: 1);
vec = std::min<unsigned>(vec, valTy.getElementType().isF16() ? 2 : 1);
// mask
numElems = tensorTy.getNumElements();
}
Expand All @@ -787,22 +783,13 @@ struct AtomicRMWOpConversion
auto vecTy = vec_ty(valueElemTy, vec);
auto retType = vec == 1 ? valueElemTy : vecTy;
SmallVector<Value> resultVals(elemsPerThread);
const bool f16v2 = vec == 2 && valueElemTy.isF16();
for (size_t i = 0; i < elemsPerThread; i += vec) {
Value rmwPtr = ptrElements[i];
// TODO: in case llMask is zero we can create only one branch for all
// elemsPerThread.
Value rmwMask = llMask ? and_(mask, maskElements[i]) : mask;

Value operand;
if (vec == 1) {
operand = valElements[i];
} else {
operand = undef(vecTy);
for (size_t ii = 0; ii < vec; ++ii)
operand =
insert_element(vecTy, operand, valElements[i + ii], i32_val(ii));
}

Value undefVal = undef(retType);
// Build blocks to bypass the atomic instruction for ~rmwMask.
auto *curBlock = rewriter.getInsertionBlock();
Expand All @@ -819,11 +806,25 @@ struct AtomicRMWOpConversion
auto maybeKind = matchAtomicOp(atomicRmwAttr);
// TODO: use rocdl.raw.buffer.atomic from ROCDL dialect to use efficient
// atomics for MI-* series of AMD GPU.
Value atom =
rewriter
.create<LLVM::AtomicRMWOp>(loc, *maybeKind, rmwPtr, operand,
atomicMemOrdering, StringRef("agent"))
.getResult();
Value atom = rewriter
.create<LLVM::AtomicRMWOp>(
loc, *maybeKind, rmwPtr, valElements[i],
atomicMemOrdering, StringRef("agent"))
.getResult();

// NV for the f16v2 case generates one packed instruction. We have to
// create two separate instructions since LLVM::AtomicRMWOp doesn't
// support this. Can be optimized out with rocdl.raw.buffer.atomic.
if (f16v2) {
Value atom2 =
rewriter
.create<LLVM::AtomicRMWOp>(
loc, *maybeKind, ptrElements[i + 1], valElements[i + 1],
atomicMemOrdering, StringRef("agent"))
.getResult();
auto tmp = insert_element(vecTy, undef(vecTy), atom, i32_val(0));
atom = insert_element(vecTy, tmp, atom2, i32_val(1)).getResult();
}
if (!tensorTy) {
if (atomicNeedsSharedMemory(op.getResult())) {
Value atomPtr =
Expand Down
Loading