saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Use CUDA's IsAllFinite kernel for ROCm

Browse source
This commit is contained in:
Jesse Benson 2020-11-18 09:41:08 -08:00 committed by Jesse Benson
parent 06ad516a5d
commit bd96f60888
5 changed files with 1 additions and 181 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
onnxruntime/core/providers/rocm/fpgeneric.cu
View file

@ -1,7 +1,7 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "hip/hip_runtime.h"
#include <hip/hip_runtime.h>
#include "core/providers/rocm/cu_inc/common.cuh"
#define TRANS_TILE_DIM 32

106
orttraining/orttraining/training_ops/rocm/math/isfinite.cc
View file

@ -1,106 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "orttraining/training_ops/rocm/math/isfinite.h"
using namespace ONNX_NAMESPACE;
using namespace onnxruntime::common;
namespace onnxruntime {
namespace rocm {
#define REGISTER_ISFINITE_KERNEL_TYPED(T) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
IsFinite, \
kMSDomain, \
1, \
T, \
kRocmExecutionProvider, \
KernelDefBuilder() \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \
.TypeConstraint("T1", DataTypeImpl::GetTensorType<bool>()), \
IsFiniteOp<T>);
template <typename TSrc>
Status IsFiniteOp<TSrc>::ComputeInternal(OpKernelContext* context) const {
typedef typename ToHipType<TSrc>::MappedType HipTSrc;
const Tensor& input = *context->Input<Tensor>(0);
Tensor& output = *context->Output(0, input.Shape());
IsFinite(
reinterpret_cast<const HipTSrc*>(input.Data<TSrc>()),
output.MutableData<bool>(), input.Shape().Size());
return Status::OK();
}
REGISTER_ISFINITE_KERNEL_TYPED(MLFloat16)
REGISTER_ISFINITE_KERNEL_TYPED(float)
REGISTER_ISFINITE_KERNEL_TYPED(double)
#define REGISTER_ISALLFINITE_KERNEL_TYPED(T) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
IsAllFinite, \
kMSDomain, \
1, \
T, \
kRocmExecutionProvider, \
KernelDefBuilder() \
.OutputMemoryType<OrtMemTypeCPUOutput>(0) \
.TypeConstraint("V", DataTypeImpl::GetTensorType<T>()) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<bool>()), \
IsAllFiniteOp<T>);
template <typename TSrc>
Status IsAllFiniteOp<TSrc>::ComputeInternal(OpKernelContext* context) const {
typedef typename ToHipType<TSrc>::MappedType TSrcCuda;
// Get Input tensor count.
const auto total_tensor_count = context->InputCount();
// Allocate GPU memory to capture the result computed by GPU kernel.
// The GPU result will be copied later to the output which locates
// on CPU memory.
IAllocatorUniquePtr<bool> deviceOutput = GetScratchBuffer<bool>(1);
HIP_RETURN_IF_ERROR(hipMemsetAsync(deviceOutput.get(), int(true), sizeof(bool)));
for (int i = 0; i < total_tensor_count; ++i) {
const auto& input = context->Input<Tensor>(i);
IsFinite(reinterpret_cast<const TSrcCuda*>(input->template Data<TSrc>()), deviceOutput.get(), input->Shape().Size());
}
// std::vector<std::vector<void*>> grouped_tensor_pointers(total_tensor_count);
// std::vector<int> tensor_sizes(total_tensor_count);
// for (int i = 0; i < total_tensor_count; ++i) {
// const auto& input = context->Input<Tensor>(i);
// grouped_tensor_pointers[i] = {const_cast<TSrc*>(input->Data<TSrc>())};
// tensor_sizes[i] = static_cast<int>(input->Shape().Size());
// IsFinite(const TSrc* input, bool* output, size_t count)
// }
// typedef IsAllFiniteFunctor<TSrcCuda> TFunctor;
// TFunctor functor;
// // Check if all values are finite and write true to deviceOutput.
// // Otherwise, false will be written.
// launch_multi_tensor_functor<1, TFunctor, bool*>(
// 2048 * 32, tensor_sizes, grouped_tensor_pointers, functor, deviceOutput.get());
// Copy GPU result in deviceOutput to CPU memory.
// Per this operator's schema, it's output is in CPU memory.
Tensor& output = *context->Output(0, {});
HIP_RETURN_IF_ERROR(
hipMemcpy(
output.MutableData<bool>(),
deviceOutput.get(),
sizeof(bool),
hipMemcpyDeviceToHost));
return Status::OK();
}
REGISTER_ISALLFINITE_KERNEL_TYPED(MLFloat16)
REGISTER_ISALLFINITE_KERNEL_TYPED(float)
REGISTER_ISALLFINITE_KERNEL_TYPED(double)
} // namespace rocm
} // namespace onnxruntime

22
orttraining/orttraining/training_ops/rocm/math/isfinite.cuh
View file

@ -1,22 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <hip/hip_fp16.h>
#include "core/providers/rocm/cu_inc/common.cuh"
#include "orttraining/training_ops/rocm/math/isfinite.h"
namespace onnxruntime {
namespace rocm {
template<typename T>
__device__ __forceinline__ bool _IsFiniteScalar(const T value) {
return isfinite(value);
}
template<>
__device__ __forceinline__ bool _IsFiniteScalar(const half value) {
return !__hisinf(value) && !__hisnan(value);
}
} // namespace rocm
} // namespace onnxruntime

49
orttraining/orttraining/training_ops/rocm/math/isfinite.h
View file

@ -1,49 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
#include "core/providers/rocm/rocm_common.h"
#include "core/providers/rocm/multi_tensor/common.cuh"
constexpr int PARALLEL_LOADS = 4;
constexpr int WARP_THREAD_COUNT = 32;
constexpr int MAX_BLOCK_COUNT = 288;
constexpr int MAX_TENSOR_COUNT = 128;
constexpr int MAX_BLOCK_THREAD_COUNT = 512;
namespace onnxruntime {
namespace rocm {
template <typename TSrc>
void IsFinite(const TSrc* input, bool* output, size_t count);
template <typename TSrc>
class IsFiniteOp final : public RocmKernel {
public:
IsFiniteOp(const OpKernelInfo& info) : RocmKernel(info) {
}
Status ComputeInternal(OpKernelContext* context) const override;
};
template <typename TSrc>
void IsFinite(const TSrc* input, bool* output, size_t N);
template <typename TSrc>
class IsAllFiniteOp final : public RocmKernel {
public:
IsAllFiniteOp(const OpKernelInfo& info) : RocmKernel(info) {
}
Status ComputeInternal(OpKernelContext* context) const override;
};
template <typename T>
struct IsAllFiniteFunctor {
void operator()(ChunkGroup<1> chunks, bool* output);
};
} // namespace rocm
} // namespace onnxruntime

3
tools/ci_build/amd_hipify.py
View file

@ -253,9 +253,6 @@ training_ops_files = [
'math/div_grad.h',
'math/div_grad_impl.cu',
'math/div_grad_impl.h',
'math/isfinite.cc',
'math/isfinite.cuh',
'math/isfinite.h',
'math/scale.cc',
'math/scale.cu',
'math/scale.h',