gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

more pass template

Browse Source
This commit is contained in:
Philippe Tillet
2023-01-06 14:26:06 -08:00
parent b16aeb6541
commit 18c7a72973
9 changed files with 92 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
include/triton/Dialect/TritonGPU/Transforms/Passes.h
View File

@@ -17,6 +17,8 @@ std::unique_ptr<Pass> createTritonGPUOptimizeLoadConvertPass();
std::unique_ptr<Pass> createTritonGPUSinkConversionsFromSharedPass();
std::unique_ptr<Pass> createTritonGPUDecomposeConversionsToDotOperandPass();
std::unique_ptr<Pass> createTritonGPUCombineOpsPass(int computeCapability = 80);
std::unique_ptr<Pass> createTritonGPUVerifier();

11
include/triton/Dialect/TritonGPU/Transforms/Passes.td
View File

@@ -96,6 +96,17 @@ def TritonGPUSinkConversionsFromShared: Pass<"tritongpu-sink-conversions-from-sh
"mlir::triton::TritonDialect"];
}
def TritonGPUDecomposeConversionsToDotOperand: Pass<"tritongpu-decompose-conversions-to-dot-operand", "mlir::ModuleOp"> {
let summary = "Decompose convert[distributed -> dotOperand] into convert[distributed -> shared -> dotOperand]";
let description = "Decomposing conversions this way makes it possible to use CSE and re-use #shared tensors";
let constructor = "mlir::createTritonGPUDecomposeConversionsToDotOperandPass()";
let dependentDialects = ["mlir::triton::gpu::TritonGPUDialect",
"mlir::triton::TritonDialect"];
}
def TritonGPUCanonicalizeLoops: Pass<"tritongpu-canonicalize-loops", "mlir::ModuleOp"> {
let summary = "canonicalize scf.ForOp ops";

36
lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp
View File

@@ -72,24 +72,24 @@ void storeDistributedToShared(Value src, Value llSrc,
Value staIdx1 = i32_val(0);
Value stride0 = dstStrides[outOrd[0]];
Value stride1 = dstStrides[outOrd[1]];
if (auto addOp = dyn_cast<LLVM::AddOp>(dynIdx0.getDefiningOp()))
if (auto cstRhs =
dyn_cast<LLVM::ConstantOp>(addOp.getRhs().getDefiningOp())) {
unsigned rhsVal =
cstRhs.getValue().cast<IntegerAttr>().getValue().getSExtValue();
unsigned key = (rhsVal / outVec) % maxPhase;
if (cache.find(key) == cache.end())
cache[key] = dynIdx0;
dynIdx0 = cache[key];
staIdx0 =
i32_val((rhsVal) / (outVec * maxPhase) * (outVec * maxPhase));
}
if (auto addOp = dyn_cast<LLVM::AddOp>(dynIdx1.getDefiningOp()))
if (auto cstRhs =
dyn_cast<LLVM::ConstantOp>(addOp.getRhs().getDefiningOp())) {
dynIdx1 = addOp.getLhs();
staIdx1 = addOp.getRhs();
}
// if (auto addOp = dyn_cast<LLVM::AddOp>(dynIdx0.getDefiningOp()))
// if (auto cstRhs =
// dyn_cast<LLVM::ConstantOp>(addOp.getRhs().getDefiningOp())) {
// unsigned rhsVal =
// cstRhs.getValue().cast<IntegerAttr>().getValue().getSExtValue();
// unsigned key = (rhsVal / outVec) % maxPhase;
// if (cache.find(key) == cache.end())
// cache[key] = dynIdx0;
// dynIdx0 = cache[key];
// staIdx0 =
// i32_val((rhsVal) / (outVec * maxPhase) * (outVec * maxPhase));
// }
// if (auto addOp = dyn_cast<LLVM::AddOp>(dynIdx1.getDefiningOp()))
// if (auto cstRhs =
// dyn_cast<LLVM::ConstantOp>(addOp.getRhs().getDefiningOp())) {
// dynIdx1 = addOp.getLhs();
// staIdx1 = addOp.getRhs();
// }
// offset along non-contiguous dimension
Value off1 = mul(dynIdx1, stride1);

1
lib/Dialect/TritonGPU/Transforms/CMakeLists.txt
View File

@@ -10,6 +10,7 @@ add_mlir_dialect_library(TritonGPUTransforms
Prefetch.cpp
OptimizeLoadConvert.cpp
SinkConversionsFromShared.cpp
DecomposeConversionsToDotOperand.cpp
TritonGPUConversion.cpp
DEPENDS

4
lib/Dialect/TritonGPU/Transforms/Combine.cpp
View File

@@ -19,7 +19,6 @@
#include <memory>
using namespace mlir;
namespace {
#include "TritonGPUCombine.inc"
@@ -483,8 +482,7 @@ public:
return op->getBlock() == cvt->getBlock() &&
!(isa<triton::ReduceOp>(op) &&
!op->getResult(0).getType().isa<RankedTensorType>()) &&
!isa<triton::gpu::ConvertLayoutOp>(op) &&
!isa<scf::YieldOp>(op);
!isa<triton::gpu::ConvertLayoutOp>(op) && !isa<scf::YieldOp>(op);
};
mlir::getForwardSlice(cvt.getResult(), &cvtSlices, filter);
if (cvtSlices.empty())

36
lib/Dialect/TritonGPU/Transforms/DecomposeConversionsToDotOperand.cpp Normal file
View File

@@ -0,0 +1,36 @@
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/SCF/SCF.h"
#include "mlir/IR/BlockAndValueMapping.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/Verifier.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/Passes.h"
#include "mlir/Transforms/RegionUtils.h"
#include "triton/Analysis/Utility.h"
#include "triton/Dialect/TritonGPU/IR/Dialect.h"
#include "triton/Dialect/TritonGPU/Transforms/Passes.h"
#include "triton/Dialect/TritonGPU/Transforms/TritonGPUConversion.h"
#define GEN_PASS_CLASSES
#include "triton/Dialect/TritonGPU/Transforms/Passes.h.inc"
using namespace mlir;
class TritonGPUDecomposeConversionsToDotOperandPass
: public TritonGPUDecomposeConversionsToDotOperandBase<
TritonGPUDecomposeConversionsToDotOperandPass> {
public:
TritonGPUDecomposeConversionsToDotOperandPass() = default;
void runOnOperation() override { return; }
};
std::unique_ptr<Pass>
mlir::createTritonGPUDecomposeConversionsToDotOperandPass() {
return std::make_unique<TritonGPUDecomposeConversionsToDotOperandPass>();
}

18
python/src/triton.cc
View File

@@ -1345,14 +1345,18 @@ void init_triton_ir(py::module &&m) {
mlir::createTritonGPUCombineOpsPass(computeCapability));
})
.def("add_tritongpu_optimize_load_convert_pass",
[](mlir::PassManager &self) {
self.addPass(mlir::createTritonGPUOptimizeLoadConvertPass());
})
[](mlir::PassManager &self) {
self.addPass(mlir::createTritonGPUOptimizeLoadConvertPass());
})
.def("add_tritongpu_sink_conversions_from_shared_pass",
[](mlir::PassManager &self) {
self.addPass(
mlir::createTritonGPUSinkConversionsFromSharedPass());
})
[](mlir::PassManager &self) {
self.addPass(mlir::createTritonGPUSinkConversionsFromSharedPass());
})
.def("add_tritongpu_decompose_conversions_to_dot_operand_pass",
[](mlir::PassManager &self) {
self.addPass(
mlir::createTritonGPUDecomposeConversionsToDotOperandPass());
})
.def("add_triton_gpu_to_llvm",
[](mlir::PassManager &self) {
self.addPass(mlir::triton::createConvertTritonGPUToLLVMPass());

2
python/triton/compiler.py
View File

@@ -906,6 +906,8 @@ def ttir_to_ttgir(mod, num_warps, num_stages, compute_capability):
pm.add_tritongpu_combine_pass(compute_capability)
pm.add_cse_pass()
# pm.add_tritongpu_optimize_load_convert_pass()
pm.add_tritongpu_sink_conversions_from_shared_pass()
pm.add_tritongpu_decompose_conversions_to_dot_operand_pass()
pm.run(mod)
return mod

14
python/tutorials/06-fused-attention.py
View File

@@ -194,8 +194,10 @@ def _bwd_kernel(
# _bwd_kernel = triton.compile("./bwd.ttgir", num_warps=8)
# _fwd_kernel = triton.compile("./fails.ptx", num_warps=4, shared=18432)
empty = torch.empty(128, device="cuda")
class _attention(torch.autograd.Function):
@staticmethod
@@ -284,8 +286,8 @@ class _attention(torch.autograd.Function):
BLOCK_DMODEL=ctx.BLOCK_DMODEL, num_warps=8,
num_stages=1,
)
print(pgm.asm["ttgir"])
exit(1)
# print(pgm.asm["ttgir"])
# exit()
return dq, dk, dv, None
@@ -327,6 +329,7 @@ def test_op(Z, H, N_CTX, D_HEAD, dtype=torch.float16):
triton.testing.assert_almost_equal(ref_dk, tri_dk)
triton.testing.assert_almost_equal(ref_dq, tri_dq)
BATCH, N_HEADS, N_CTX, D_HEAD = 4, 48, 4096, 64
# vary seq length for fixed head and batch=4
configs = [triton.testing.Benchmark(
@@ -358,8 +361,8 @@ def bench_flash_attention(BATCH, H, N_CTX, D_HEAD, mode, provider, dtype=torch.f
do = torch.randn_like(o)
fn = lambda: o.backward(do, retain_graph=True)
ms = triton.testing.do_bench(fn, percentiles=None, warmup=warmup, rep=rep)
flops_per_matmul = 2.*BATCH*H*N_CTX*N_CTX*D_HEAD*0.5
total_flops = 2*flops_per_matmul
flops_per_matmul = 2. * BATCH * H * N_CTX * N_CTX * D_HEAD * 0.5
total_flops = 2 * flops_per_matmul
# print(total_flops/ms*1e-9)
print(ms)
return ms
@@ -376,4 +379,5 @@ def bench_flash_attention(BATCH, H, N_CTX, D_HEAD, mode, provider, dtype=torch.f
ms = triton.testing.do_bench(fn, percentiles=None, warmup=warmup, rep=rep)
return ms
bench_flash_attention.run(save_path='.', print_data=True)
# bench_flash_attention.run(save_path='.', print_data=True)