Add AMDGPU dialect scaffolding #4685
Conversation
| add_subdirectory(TritonAMDGPUToLLVM) | ||
| add_subdirectory(TritonAMDGPUTransforms) | ||
| add_subdirectory(Dialect) |
There was a problem hiding this comment.
You wanna put Dialect as the first given others may depends on it.
| #include "mlir/Dialect/LLVMIR/LLVMDialect.h" | ||
| #include "mlir/Pass/Pass.h" | ||
|
|
||
| namespace mlir { |
There was a problem hiding this comment.
Nit: namespace mlir::trition.
|
|
||
| def ConvertAMDGPUToLLVM : Pass<"convert-amd-gpu-to-llvm", "mlir::ModuleOp"> { | ||
| let summary = "Convert AMDGPU to LLVM"; | ||
| let description = [{ |
There was a problem hiding this comment.
Can we put some description here?
| let cppNamespace = "::mlir::triton::amdgpu"; | ||
|
|
||
| let description = [{ | ||
| AMDGPU Dialect. |
There was a problem hiding this comment.
let summary = "Triton AMDGPU dialect".
let description = [{
This dialect hosts operations used for lowering patterns in the Triton AMD backend.
}];
| */ | ||
|
|
||
|
|
||
| #ifndef AMDGPU_OPS |
There was a problem hiding this comment.
Can we prefix these #ifdef in TD files with TRITON_ to avoid potential collision with other sources? Upstream has an AMDGPU dialect too..
Also for other TD files in this patch.
| } | ||
| } // namespace mlir::triton::AMD | ||
|
|
||
| class ConvertAMDGPUToLLVM |
There was a problem hiding this comment.
Nit: can use struct given no private fields anyway.
|
|
||
| using namespace mlir; | ||
| using namespace mlir::triton; | ||
| using namespace mlir::triton::gpu; |
There was a problem hiding this comment.
Not a fan of just using all namespaces. Can we drop all except the first two? We can consider add them if necessary later. But don't want to just enable all from the get-go.
| mlir::LowerToLLVMOptions option(context); | ||
|
|
||
| TritonGPUToLLVMTypeConverter typeConverter(context, option); | ||
| ModuleAxisInfoAnalysis axisInfoAnalysis(mod); |
There was a problem hiding this comment.
Don't think we need this right now? Can we delete and add later when really needed?
| public: | ||
| explicit AMDDialectLLVMConversionTarget(MLIRContext &ctx) | ||
| : ConversionTarget(ctx) { | ||
| addLegalDialect<LLVM::LLVMDialect>(); |
There was a problem hiding this comment.
You can pass multiple dialect into the same addLegalDialect<...> call.
| addLegalDialect<mlir::scf::SCFDialect>(); | ||
| addLegalDialect<triton::TritonDialect>(); | ||
| addLegalDialect<triton::gpu::TritonGPUDialect>(); | ||
| addIllegalDialect<triton::nvidia_gpu::TritonNvidiaGPUDialect>(); |
There was a problem hiding this comment.
For illegal ones, maybe use markUnknownOpDynamicallyLegal to return false for all of them? This way we are a bit future proof so that we don't need to remember updating this after dialect changes.
third_party/amd/python/triton_amd.cc
Outdated
| @@ -64,6 +66,8 @@ void init_triton_amd_passes_ttgpuir(py::module &&m) { | |||
| mlir::createTritonAMDGPUStreamPipelinePass); | |||
| ADD_PASS_WRAPPER_1("add_stream_pipelinev2", | |||
| mlir::createTritonAMDGPUStreamPipelineV2Pass, int); | |||
| ADD_PASS_WRAPPER_0("add_amdgpu_to_llvm", | |||
| mlir::triton::createConvertAMDGPUToLLVMPass); | |||
There was a problem hiding this comment.
why do we need a separate pass for AMDGPU to LLVM? Ideally this is done with extra pattern on top of TritonGPU to llvm as breaking up the conversion to llvm into multiple pass has downsides. It means we get an intermediate IR that mixed block and simt semantic and it prevents pattern matching when lowering.
Yup at TritonGPU level. I just realized that the initial commit message isn't providing enough anticipated usage of this dialect. Basically it's meant to complement TritonGPU in the lowering process for AMD side like what we have for NVIDIA ones. Example ops we are considering right now (not definitely confirmed we will have them):
|
That's what I was expecting and I think it makes sense. One think to consider for the naming, on Nivida side the analogue dialect is called |
|
Yeah that makes sense--we want to be consistent regarding naming there. Better to have that done at the beginning. |
oplavsic
force-pushed
the
amd_dialect_load_wip
branch
from
e4084c5
to
c622c1e
Compare
| let cppNamespace = "::mlir::triton::amdgpu"; | ||
|
|
||
| let description = [{ | ||
| TritonAMDGPU Dialect is used to support AMD specific instructions on TritonGPU level abstraction. |
There was a problem hiding this comment.
.. hosts AMD specific ops at TritonGPU abstraction level.
There was a problem hiding this comment.
Done. Thanks for the review!
oplavsic
force-pushed
the
amd_dialect_load_wip
branch
from
c622c1e
to
c1d1a3b
Compare
This PR adds a scaffolding for defining and lowering AMD specific dialect to LLVM.