wehub-resource-sync
57 lines
1.7 KiB
C++
57 lines
1.7 KiB
C++
// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
// The file has been adapted from pytorch project
|
|
// Licensed under BSD-style license -
|
|
// https://github.com/pytorch/pytorch/blob/main/LICENSE
|
|
|
|
#pragma once
|
|
|
|
#include <cstdint>
|
|
|
|
namespace at {
|
|
|
|
struct PhiloxCudaState {
|
|
PhiloxCudaState() = default;
|
|
// Called if graph capture is not underway
|
|
PhiloxCudaState(uint64_t seed, uint64_t offset) {
|
|
seed_.val = seed;
|
|
offset_.val = offset;
|
|
}
|
|
// Called if graph capture is underway
|
|
PhiloxCudaState(int64_t* seed,
|
|
int64_t* offset_extragraph,
|
|
uint64_t offset_intragraph) {
|
|
seed_.ptr = seed;
|
|
offset_.ptr = offset_extragraph;
|
|
offset_intragraph_ = offset_intragraph;
|
|
captured_ = true;
|
|
}
|
|
|
|
// Public members, directly accessible by at::cuda::philox::unpack.
|
|
// If we made them private with getters/setters, the getters/setters
|
|
// would have to be __device__, and we can't declare __device__ in ATen.
|
|
union Payload {
|
|
uint64_t val;
|
|
int64_t* ptr;
|
|
};
|
|
|
|
Payload seed_{};
|
|
Payload offset_{};
|
|
uint64_t offset_intragraph_ = 0;
|
|
bool captured_ = false;
|
|
};
|
|
|
|
} // namespace at
|