02ebf24d35
The purpose of this PR is analyzing shared memory aliases so that we can fix memory allocation bugs and save memory allocations in triton code involving complex control flows. Changes to memory bar and allocation are on the way. Co-authored-by: Philippe Tillet <phil@openai.com>
43 lines
1.3 KiB
C++
43 lines
1.3 KiB
C++
#include "triton/Dialect/Triton/IR/Dialect.h"
|
|
#include "triton/Dialect/TritonGPU/IR/Dialect.h"
|
|
|
|
#include "triton/Dialect/Triton/Transforms/Passes.h"
|
|
#include "triton/Dialect/TritonGPU/Transforms/Passes.h"
|
|
|
|
#include "triton/Conversion/Passes.h"
|
|
|
|
#include "mlir/IR/Dialect.h"
|
|
#include "mlir/InitAllPasses.h"
|
|
#include "mlir/Support/MlirOptMain.h"
|
|
|
|
namespace mlir {
|
|
namespace test {
|
|
void registerTestAliasPass();
|
|
void registerTestAlignmentPass();
|
|
void registerTestAllocationPass();
|
|
void registerTestMembarPass();
|
|
} // namespace test
|
|
} // namespace mlir
|
|
|
|
int main(int argc, char **argv) {
|
|
mlir::registerAllPasses();
|
|
mlir::registerTritonPasses();
|
|
mlir::registerTritonGPUPasses();
|
|
mlir::test::registerTestAliasPass();
|
|
mlir::test::registerTestAlignmentPass();
|
|
mlir::test::registerTestAllocationPass();
|
|
mlir::test::registerTestMembarPass();
|
|
mlir::triton::registerConvertTritonToTritonGPUPass();
|
|
mlir::triton::registerConvertTritonGPUToLLVMPass();
|
|
|
|
// TODO: register Triton & TritonGPU passes
|
|
mlir::DialectRegistry registry;
|
|
registry
|
|
.insert<mlir::triton::TritonDialect, mlir::triton::gpu::TritonGPUDialect,
|
|
mlir::arith::ArithmeticDialect, mlir::StandardOpsDialect,
|
|
mlir::scf::SCFDialect>();
|
|
|
|
return mlir::asMainReturnCode(mlir::MlirOptMain(
|
|
argc, argv, "Triton (GPU) optimizer driver\n", registry));
|
|
}
|