-
Notifications
You must be signed in to change notification settings - Fork 2.3k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[lang] Instantiate a runtime ArgPack object when a python ArgPack is …
…created ghstack-source-id: d1dc8a846ff1202b84fc47c270d0e88e7dabeca3 Pull Request resolved: #8241
- Loading branch information
Showing
13 changed files
with
275 additions
and
11 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,49 @@ | ||
#include <numeric> | ||
|
||
#include "taichi/program/argpack.h" | ||
#include "taichi/program/program.h" | ||
|
||
#ifdef TI_WITH_LLVM | ||
#include "taichi/runtime/llvm/llvm_context.h" | ||
#include "taichi/runtime/program_impls/llvm/llvm_program.h" | ||
#endif | ||
|
||
namespace taichi::lang { | ||
|
||
ArgPack::ArgPack(Program *prog, | ||
const DataType type) | ||
: prog_(prog) { | ||
|
||
auto* old_type = type->get_type()->as<ArgPackType>(); | ||
auto [argpack_type, alloc_size] = prog->get_argpack_type_with_data_layout( | ||
old_type, old_type->get_layout()); | ||
dtype = DataType(argpack_type); | ||
argpack_alloc_ = prog->allocate_memory_on_device(alloc_size, | ||
prog->result_buffer); | ||
} | ||
|
||
ArgPack::~ArgPack() { | ||
if (prog_) { | ||
argpack_alloc_.device->dealloc_memory(argpack_alloc_); | ||
} | ||
} | ||
|
||
intptr_t ArgPack::get_device_allocation_ptr_as_int() const { | ||
// taichi's own argpack's ptr points to its |DeviceAllocation| on the | ||
// specified device. | ||
return reinterpret_cast<intptr_t>(&argpack_alloc_); | ||
} | ||
|
||
DeviceAllocation ArgPack::get_device_allocation() const { | ||
return argpack_alloc_; | ||
} | ||
|
||
std::size_t ArgPack::get_nelement() const { | ||
return dtype->as<ArgPackType>()->elements().size(); | ||
} | ||
|
||
DataType ArgPack::get_data_type() const { | ||
return dtype; | ||
} | ||
|
||
} // namespace taichi::lang |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,35 @@ | ||
#pragma once | ||
|
||
#include <cstdint> | ||
#include <vector> | ||
|
||
#include "taichi/inc/constants.h" | ||
#include "taichi/ir/type_utils.h" | ||
#include "taichi/rhi/device.h" | ||
|
||
namespace taichi::lang { | ||
|
||
class Program; | ||
|
||
class TI_DLL_EXPORT ArgPack { | ||
public: | ||
/* Constructs a ArgPack managed by Program. | ||
* Memory allocation and deallocation is handled by Program. | ||
*/ | ||
explicit ArgPack(Program *prog, | ||
const DataType type); | ||
|
||
DeviceAllocation argpack_alloc_{kDeviceNullAllocation}; | ||
DataType dtype; | ||
|
||
DataType get_data_type() const; | ||
intptr_t get_device_allocation_ptr_as_int() const; | ||
DeviceAllocation get_device_allocation() const; | ||
std::size_t get_nelement() const; | ||
|
||
~ArgPack(); | ||
|
||
private: | ||
Program *prog_{nullptr}; | ||
}; | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.