Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix a compilation error when compiling with half and cuTENSOR #2450

Merged
merged 1 commit into from
Jun 6, 2024
Merged

Fix a compilation error when compiling with half and cuTENSOR #2450

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
235 changes: 120 additions & 115 deletions src/layers/transform/multidim_reduction.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,139 +148,144 @@ void multidim_reduction_layer<TensorDataType, Layout, Device>::fp_compute()
// No model parallelism nor CPU
LBANN_ASSERT(Layout == data_layout::DATA_PARALLEL);
LBANN_ASSERT(Device == El::Device::GPU);
if constexpr (Device == El::Device::GPU)
{
// Constants
const auto one = El::TypeTraits<TensorDataType>::One();
const auto zero = El::TypeTraits<TensorDataType>::Zero();

// Constants
const auto one = El::TypeTraits<TensorDataType>::One();
const auto zero = El::TypeTraits<TensorDataType>::Zero();
// Data matrices
using LocalMat = El::Matrix<TensorDataType, El::Device::GPU>;
LocalMat const& input =
static_cast<LocalMat const&>(this->get_prev_activations().LockedMatrix());
LocalMat& output = static_cast<LocalMat&>(this->get_activations().Matrix());

// Data matrices
using LocalMat = El::Matrix<TensorDataType, El::Device::GPU>;
LocalMat const& input =
static_cast<LocalMat const&>(this->get_prev_activations().LockedMatrix());
LocalMat& output = static_cast<LocalMat&>(this->get_activations().Matrix());

// Determine reduction operator
cutensorOperator_t cutensor_reduce_op;
switch (m_mode) {
case multidim_reduction_mode::SUM:
cutensor_reduce_op = CUTENSOR_OP_ADD;
break;
// The following reduction modes are enabled only in forward prop
case multidim_reduction_mode::PRODUCT:
cutensor_reduce_op = CUTENSOR_OP_MUL;
break;
case multidim_reduction_mode::MAX:
cutensor_reduce_op = CUTENSOR_OP_MAX;
break;
case multidim_reduction_mode::MIN:
cutensor_reduce_op = CUTENSOR_OP_MIN;
break;
default:
LBANN_ERROR("invalid reduction mode");
}
// Determine reduction operator
cutensorOperator_t cutensor_reduce_op;
switch (m_mode) {
case multidim_reduction_mode::SUM:
cutensor_reduce_op = CUTENSOR_OP_ADD;
break;
// The following reduction modes are enabled only in forward prop
case multidim_reduction_mode::PRODUCT:
cutensor_reduce_op = CUTENSOR_OP_MUL;
break;
case multidim_reduction_mode::MAX:
cutensor_reduce_op = CUTENSOR_OP_MAX;
break;
case multidim_reduction_mode::MIN:
cutensor_reduce_op = CUTENSOR_OP_MIN;
break;
default:
LBANN_ERROR("invalid reduction mode");
}

RowMajorDims<int64_t> input_dims(this->get_input_dims());
RowMajorDims<int64_t> output_dims(this->get_output_dims());
auto const& input_modes = m_input_modes;
auto const& output_modes = m_output_modes;
RowMajorDims<int64_t> input_dims(this->get_input_dims());
RowMajorDims<int64_t> output_dims(this->get_output_dims());
auto const& input_modes = m_input_modes;
auto const& output_modes = m_output_modes;

auto input_desc = get_descriptor(input, input_dims);
auto output_desc = get_descriptor(output, output_dims);
auto input_desc = get_descriptor(input, input_dims);
auto output_desc = get_descriptor(output, output_dims);

// Create workspace buffers
LocalMat workspace;
workspace.SetMemoryMode(1);
auto handle = get_handle_ptr();
uint64_t wspsize = 0;
CHECK_CUTENSOR(cutensorReductionGetWorkspaceSize(handle,
input.LockedBuffer(),
&input_desc,
input_modes.data(),
output.LockedBuffer(),
&output_desc,
output_modes.data(),
output.LockedBuffer(),
&output_desc,
output_modes.data(),
cutensor_reduce_op,
CUTENSOR_COMPUTE_32F,
&wspsize));
workspace.Resize(wspsize, 1);
// Create workspace buffers
LocalMat workspace;
workspace.SetMemoryMode(1);
auto handle = get_handle_ptr();
uint64_t wspsize = 0;
CHECK_CUTENSOR(cutensorReductionGetWorkspaceSize(handle,
input.LockedBuffer(),
&input_desc,
input_modes.data(),
output.LockedBuffer(),
&output_desc,
output_modes.data(),
output.LockedBuffer(),
&output_desc,
output_modes.data(),
cutensor_reduce_op,
CUTENSOR_COMPUTE_32F,
&wspsize));
workspace.Resize(wspsize, 1);

auto multisync = El::MakeMultiSync(El::SyncInfoFromMatrix(output),
El::SyncInfoFromMatrix(input));
auto multisync = El::MakeMultiSync(El::SyncInfoFromMatrix(output),
El::SyncInfoFromMatrix(input));

// Compute reduction locally
CHECK_CUTENSOR(cutensorReduction(
handle,
&one,
input.LockedBuffer(),
&input_desc,
input_modes.data(),
&zero,
output.LockedBuffer(),
&output_desc,
output_modes.data(),
output.Buffer(),
&output_desc,
output_modes.data(),
cutensor_reduce_op,
CUTENSOR_COMPUTE_32F,
workspace.Buffer(),
wspsize,
static_cast<El::SyncInfo<El::Device::GPU>>(multisync).Stream()));
// Compute reduction locally
CHECK_CUTENSOR(cutensorReduction(
handle,
&one,
input.LockedBuffer(),
&input_desc,
input_modes.data(),
&zero,
output.LockedBuffer(),
&output_desc,
output_modes.data(),
output.Buffer(),
&output_desc,
output_modes.data(),
cutensor_reduce_op,
CUTENSOR_COMPUTE_32F,
workspace.Buffer(),
wspsize,
static_cast<El::SyncInfo<El::Device::GPU>>(multisync).Stream()));
}
#endif
}

template <typename TensorDataType, data_layout Layout, El::Device Device>
void multidim_reduction_layer<TensorDataType, Layout, Device>::bp_compute()
{
#ifdef LBANN_HAS_CUTENSOR
// Constants
const auto one = El::TypeTraits<TensorDataType>::One();
const auto zero = El::TypeTraits<TensorDataType>::Zero();
if constexpr (Device == El::Device::GPU)
{
// Constants
const auto one = El::TypeTraits<TensorDataType>::One();
const auto zero = El::TypeTraits<TensorDataType>::Zero();

// Data matrices
using LocalMat = El::Matrix<TensorDataType, El::Device::GPU>;
LocalMat const& input =
static_cast<LocalMat const&>(this->get_prev_error_signals().LockedMatrix());
LocalMat& output = static_cast<LocalMat&>(this->get_error_signals().Matrix());
// Data matrices
using LocalMat = El::Matrix<TensorDataType, El::Device::GPU>;
LocalMat const& input =
static_cast<LocalMat const&>(this->get_prev_error_signals().LockedMatrix());
LocalMat& output = static_cast<LocalMat&>(this->get_error_signals().Matrix());

RowMajorDims<int64_t> input_dims(this->get_output_dims());
RowMajorDims<int64_t> output_dims(this->get_input_dims());
auto const& input_modes = m_output_modes;
auto const& output_modes = m_input_modes;
RowMajorDims<int64_t> input_dims(this->get_output_dims());
RowMajorDims<int64_t> output_dims(this->get_input_dims());
auto const& input_modes = m_output_modes;
auto const& output_modes = m_input_modes;

auto multisync = El::MakeMultiSync(El::SyncInfoFromMatrix(output),
El::SyncInfoFromMatrix(input));
auto handle = get_handle_ptr();
auto multisync = El::MakeMultiSync(El::SyncInfoFromMatrix(output),
El::SyncInfoFromMatrix(input));
auto handle = get_handle_ptr();

// Determine reduction operator
switch (m_mode) {
case multidim_reduction_mode::SUM: {
// Broadcast dimensions
auto input_desc = get_descriptor(input, input_dims);
auto output_desc = get_descriptor(output, output_dims);
CHECK_CUTENSOR(cutensorElementwiseBinary(
handle,
&one,
input.LockedBuffer(),
&input_desc,
input_modes.data(),
&zero,
output.LockedBuffer(),
&output_desc,
output_modes.data(),
output.Buffer(),
&output_desc,
output_modes.data(),
CUTENSOR_OP_ADD,
CUDATypeT<TensorDataType>::value,
static_cast<El::SyncInfo<El::Device::GPU>>(multisync).Stream()));
break;
}
default:
LBANN_ERROR("invalid reduction mode, only sum is supported for training");
// Determine reduction operator
switch (m_mode) {
case multidim_reduction_mode::SUM: {
// Broadcast dimensions
auto input_desc = get_descriptor(input, input_dims);
auto output_desc = get_descriptor(output, output_dims);
CHECK_CUTENSOR(cutensorElementwiseBinary(
handle,
&one,
input.LockedBuffer(),
&input_desc,
input_modes.data(),
&zero,
output.LockedBuffer(),
&output_desc,
output_modes.data(),
output.Buffer(),
&output_desc,
output_modes.data(),
CUTENSOR_OP_ADD,
CUDATypeT<TensorDataType>::value,
static_cast<El::SyncInfo<El::Device::GPU>>(multisync).Stream()));
break;
}
default:
LBANN_ERROR("invalid reduction mode, only sum is supported for training");
}
}
#endif
}
Expand Down
Loading