Skip to content

[STABLE ABI] Eliminate C10 and ATen dependencies. #4137

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

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

[STABLE ABI] Eliminate C10 and ATen dependencies. #4137

Show file tree
Hide file tree
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
5 changes: 2 additions & 3 deletions src/libtorchaudio/rnnt/cpu/cpu_kernels.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,7 @@
#include <libtorchaudio/rnnt/cpu/math.h>
#include <libtorchaudio/rnnt/options.h>
#include <libtorchaudio/rnnt/types.h>

#include <c10/util/Logging.h>
#include <torch/headeronly/util/Exception.h>

#include <cstring>
#include <limits>
Expand Down Expand Up @@ -50,7 +49,7 @@ class TensorView {
}

DTYPE& operator()(const std::vector<int>& indices) {
TORCH_CHECK_EQ(indices.size(), dims_.size());
STD_TORCH_CHECK(indices.size() == dims_.size());
int index = indices.back();
for (int i = indices.size() - 2; i >= 0; --i) {
index += indices[i] * strides_[i];
Expand Down
6 changes: 3 additions & 3 deletions src/libtorchaudio/rnnt/cpu/cpu_transducer.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ status_t Compute(
DTYPE* gradients = nullptr) {
const Options& options = workspace.GetOptions();

TORCH_CHECK_EQ(options.device_, CPU);
STD_TORCH_CHECK(options.device_ == CPU);

const int& B = options.batchSize_;
const int& maxT = options.maxSrcLen_;
Expand Down Expand Up @@ -91,7 +91,7 @@ status_t ComputeAlphas(
DTYPE* alphas) {
const Options& options = workspace.GetOptions();

TORCH_CHECK_EQ(options.device_, CPU);
STD_TORCH_CHECK(options.device_ == CPU);

const int& B = options.batchSize_;
const int& maxT = options.maxSrcLen_;
Expand Down Expand Up @@ -140,7 +140,7 @@ status_t ComputeBetas(
DTYPE* betas) {
const Options& options = workspace.GetOptions();

TORCH_CHECK_EQ(options.device_, CPU);
STD_TORCH_CHECK(options.device_ == CPU);

const int& B = options.batchSize_;
const int& maxT = options.maxSrcLen_;
Expand Down
44 changes: 17 additions & 27 deletions src/libtorchaudio/rnnt/gpu/compute.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,10 @@

#include <c10/cuda/CUDAException.h>
#include <c10/cuda/CUDAStream.h>

#include <torch/csrc/stable/library.h>
#include <torch/csrc/stable/ops.h>
#include <torch/headeronly/core/Dispatch_v2.h>

namespace torchaudio {
namespace rnnt {
Expand Down Expand Up @@ -117,33 +119,21 @@ std::tuple<Tensor, Tensor> compute(
/*int_data=*/reinterpret_cast<int*>(int_workspace.data_ptr()),
/*int_size=*/int_workspace.numel());

switch (logits.scalar_type()) {
case ScalarType::Float: {
Compute</*DTYPE=*/float, /*CAST_DTYPE=*/float>(
/*workspace=*/workspace,
/*logits=*/reinterpret_cast<float*>(logits.data_ptr()),
/*targets=*/reinterpret_cast<int*>(targets.data_ptr()),
/*srcLengths=*/reinterpret_cast<int*>(logit_lengths.data_ptr()),
/*tgtLengths=*/reinterpret_cast<int*>(target_lengths.data_ptr()),
/*costs=*/reinterpret_cast<float*>(costs.data_ptr()),
/*gradients=*/reinterpret_cast<float*>(gradients.data_ptr()));
break;
}
case ScalarType::Half: {
Compute</*DTYPE=*/c10::Half, /*CAST_DTYPE=*/float>(
/*workspace=*/workspace,
/*logits=*/reinterpret_cast<c10::Half*>(logits.data_ptr()),
/*targets=*/reinterpret_cast<int*>(targets.data_ptr()),
/*srcLengths=*/reinterpret_cast<int*>(logit_lengths.data_ptr()),
/*tgtLengths=*/reinterpret_cast<int*>(target_lengths.data_ptr()),
/*costs=*/reinterpret_cast<c10::Half*>(costs.data_ptr()),
/*gradients=*/reinterpret_cast<c10::Half*>(gradients.data_ptr()));
break;
}
default: {
STD_TORCH_CHECK(false, "unreachable");
}
};
THO_DISPATCH_V2(
logits.scalar_type(),
"rnnt:compute",
AT_WRAP([&] {
(Compute</*DTYPE=*/scalar_t, /*CAST_DTYPE=*/float>(
/*workspace=*/workspace,
/*logits=*/reinterpret_cast<scalar_t*>(logits.data_ptr()),
/*targets=*/reinterpret_cast<int*>(targets.data_ptr()),
/*srcLengths=*/reinterpret_cast<int*>(logit_lengths.data_ptr()),
/*tgtLengths=*/reinterpret_cast<int*>(target_lengths.data_ptr()),
/*costs=*/reinterpret_cast<scalar_t*>(costs.data_ptr()),
/*gradients=*/reinterpret_cast<scalar_t*>(gradients.data_ptr())));
}),
ScalarType::Float,
ScalarType::Half);

return std::make_tuple(costs, gradients);
}
Expand Down
2 changes: 1 addition & 1 deletion src/libtorchaudio/rnnt/gpu/half.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
#pragma once

#ifdef USE_C10_HALF
#include "c10/util/Half.h"
#include <torch/headeronly/util/Half.h>
#endif // USE_C10_HALF

#include <libtorchaudio/rnnt/macros.h>
Expand Down
13 changes: 6 additions & 7 deletions src/libtorchaudio/rnnt/workspace.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,7 @@
#include <vector>

#include <libtorchaudio/rnnt/options.h>

#include <c10/util/Logging.h>
#include <torch/headeronly/util/Exception.h>

namespace torchaudio {
namespace rnnt {
Expand All @@ -29,7 +28,7 @@ class DtypeWorkspace {
~DtypeWorkspace() {}

static int ComputeSizeFromOptions(const Options& options) {
TORCH_CHECK_NE(options.device_, UNDEFINED);
STD_TORCH_CHECK(options.device_ != UNDEFINED);
return ComputeSizeForDenominators(options) +
ComputeSizeForLogProbs(options) + ComputeSizeForAlphas(options) +
ComputeSizeForBetas(options);
Expand All @@ -38,7 +37,7 @@ class DtypeWorkspace {
void Free();
void Reset(const Options& options, DTYPE* data, int size) {
int needed_size = ComputeSizeFromOptions(options);
TORCH_CHECK_LE(needed_size, size);
STD_TORCH_CHECK(needed_size <= size);
options_ = options;
data_ = data;
size_ = size;
Expand Down Expand Up @@ -100,7 +99,7 @@ class IntWorkspace {

void Reset(const Options& options, int* data, int size) {
int needed_size = ComputeSizeFromOptions(options);
TORCH_CHECK_LE(needed_size, size);
STD_TORCH_CHECK(needed_size <= size);
options_ = options;
data_ = data;
size_ = size;
Expand All @@ -111,11 +110,11 @@ class IntWorkspace {
}

int* GetPointerToAlphaCounters() const {
TORCH_CHECK_EQ(options_.device_, GPU);
STD_TORCH_CHECK(options_.device_ == GPU);
return data_;
}
int* GetPointerToBetaCounters() const {
TORCH_CHECK_EQ(options_.device_, GPU);
STD_TORCH_CHECK(options_.device_ == GPU);
return GetPointerToAlphaCounters() + ComputeSizeForAlphaCounters(options_);
}

Expand Down
2 changes: 0 additions & 2 deletions src/libtorchaudio/utils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,4 @@
#include <ATen/DynamicLibrary.h>
#include <libtorchaudio/utils.h>
#include <torch/csrc/stable/tensor.h>

#ifdef USE_CUDA
#include <cuda.h>
Expand Down