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

cast to amd device when as_nvidia shows up

Browse Source
This commit is contained in:
Michael Melesse
2022-10-26 18:12:18 +00:00
parent 4464dfcc18
commit abe0d3e1b1
2 changed files with 22 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
include/triton/codegen/target.h
View File

@@ -36,6 +36,7 @@ namespace triton{
namespace codegen{ namespace codegen{
class nvidia_cu_target; class nvidia_cu_target;
class amd_cl_target;
class target { class target {
public: public:
@@ -49,7 +50,12 @@ public:
virtual Value* get_block_id(Module *module, Builder& builder, unsigned ax) = 0; virtual Value* get_block_id(Module *module, Builder& builder, unsigned ax) = 0;
virtual Value* get_num_blocks(Module *module, Builder& builder, unsigned ax) = 0; virtual Value* get_num_blocks(Module *module, Builder& builder, unsigned ax) = 0;
virtual unsigned guaranteed_alignment() = 0; virtual unsigned guaranteed_alignment() = 0;
#ifdef USE_ROCM
amd_cl_target* as_nvidia();
amd_cl_target* as_amd();
#else
nvidia_cu_target* as_nvidia(); nvidia_cu_target* as_nvidia();
#endif
bool is_gpu() const; bool is_gpu() const;
private: private:
@@ -67,6 +73,7 @@ public:
Value* get_block_id(Module *module, Builder& builder, unsigned ax); Value* get_block_id(Module *module, Builder& builder, unsigned ax);
Value* get_num_blocks(Module *module, Builder& builder, unsigned ax); Value* get_num_blocks(Module *module, Builder& builder, unsigned ax);
unsigned guaranteed_alignment() { return 16; } unsigned guaranteed_alignment() { return 16; }
int sm() { return 0; } // treat as if old CUDA device
}; };
class nvidia_cu_target: public target { class nvidia_cu_target: public target {

18
lib/codegen/target.cc
View File

@@ -15,10 +15,22 @@ namespace codegen{
// base // base
#ifdef USE_ROCM
nvidia_cu_target* target::as_nvidia() { amd_cl_target *target::as_amd()
return dynamic_cast<nvidia_cu_target*>(this); {
return dynamic_cast<amd_cl_target *>(this);
} }
amd_cl_target *target::as_nvidia()
{
return this->as_amd();
}
#else
// causes segfault on ROCM
nvidia_cu_target *target::as_nvidia()
{
return dynamic_cast<nvidia_cu_target *>(this);
}
#endif
bool target::is_gpu() const { bool target::is_gpu() const {
return is_gpu_; return is_gpu_;