mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
6d24f8fe21
Summary: **Summary:** This PR contains the infrastructure of a new CUDA fuser. This CUDA fuser is based on many of the same principles of TensorExpressions and Halide, however the implementation is ground up. The fusion pass itself is similar to the default CUDA fuser, however, it has undergone some refactoring and is using the new code generation infrastructure. For those who are interested in how the code generation in this PR works, I would recommend reviewing _test/cpp/jit/test_gpu_fusion.cpp_ as well as the long comment section at the beginning of _torch/csrc/jit/codegen/cuda/transform_replay.h_ One of the largest differences between our approach and that of TVM/Halide, is the concept of "TensorView". TensorView from a high level should be thought of similarly to how we think of working with Tensors in PyTorch. It's an N-D object which can undergo transformations that change its dimensionality. Dimensionality changes are done through the operations split/merge/reorder/computeAt. These transformations are similar to split/fuse/reorder/compute_at of TVM, they modify how a tensor is iterated over to generate GPU code. Interestingly, in our scheme these transformations are applied to tensors and only impact how that tensor is generated. **Warning:** This PR is purposefully not feature complete with the current fuser. We wanted to separate out the infrastructure from the fusion capabilities. Once in, smaller incremental PRs will be submitted to expand capabilities of the fuser. **Short term goals:** Parity with current CUDA fuser (including performance): - Dynamic shapes (no recompilation) - Implicit handling of braodcast (broadcasted tensors are treated as tensors of the braodcasted size in the generated code) - Dropout **Mid-term goals:** - Transposes fused with pointwise operations where transpose involves only 2 axes (across the fused operation). - 1-D reductions fused with pointwise operations Pull Request resolved: https://github.com/pytorch/pytorch/pull/34785 Reviewed By: ZolotukhinM Differential Revision: D20650977 Pulled By: soumith fbshipit-source-id: ee39c95a880e1b9822e874ed4cc180971572bf63
104 lines
2.8 KiB
C++
104 lines
2.8 KiB
C++
#pragma once
|
|
|
|
#include <c10/core/Device.h>
|
|
#include <c10/core/DeviceType.h>
|
|
#include <c10/util/Optional.h>
|
|
#include <torch/csrc/WindowsTorchApiMacro.h> // TORCH_CUDA_API
|
|
|
|
#include <aten/src/ATen/core/jit_type.h>
|
|
#include <torch/csrc/jit/ir/ir.h>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
namespace fuser {
|
|
|
|
using RankType = std::vector<int64_t>::size_type;
|
|
|
|
/*
|
|
* Functions for printing ATen IR
|
|
*/
|
|
|
|
TORCH_CUDA_API void printScalar(std::ostream& stream, const Value* const value);
|
|
|
|
TORCH_CUDA_API void printStrides(
|
|
std::ostream& stream,
|
|
const c10::VaryingStrides& strides);
|
|
|
|
TORCH_CUDA_API void printSizes(
|
|
std::ostream& stream,
|
|
const c10::VaryingShape& sizes);
|
|
|
|
TORCH_CUDA_API void printCompleteTensor(
|
|
std::ostream& stream,
|
|
const std::shared_ptr<c10::TensorType>& tensor);
|
|
|
|
TORCH_CUDA_API void printValue(std::ostream& stream, const Value* const value);
|
|
|
|
/*
|
|
* Functions for acquiring devices and device types from ATen IR nodes
|
|
*/
|
|
|
|
// Warning: assumes all fusion outputs are complete tensors
|
|
TORCH_CUDA_API c10::Device getFusionDevice(const Node* const node);
|
|
|
|
TORCH_CUDA_API c10::DeviceType getFusionDeviceType(const Node* const node);
|
|
|
|
/*
|
|
* Functions for obtaining parts of complete tensors
|
|
*/
|
|
|
|
TORCH_CUDA_API c10::DeviceType getDeviceType(
|
|
const std::shared_ptr<c10::TensorType>& tensor);
|
|
|
|
TORCH_CUDA_API std::vector<int64_t> extractStrides(
|
|
const std::shared_ptr<c10::TensorType>& tensor);
|
|
|
|
TORCH_CUDA_API std::vector<int64_t> extractSizes(
|
|
const std::shared_ptr<c10::TensorType>& tensor);
|
|
|
|
TORCH_CUDA_API size_t getRank(const std::shared_ptr<c10::TensorType>& tensor);
|
|
|
|
TORCH_CUDA_API size_t getNumel(const std::shared_ptr<c10::TensorType>& tensor);
|
|
|
|
/*
|
|
* Functions for working with scalar Values
|
|
*/
|
|
|
|
TORCH_CUDA_API bool isScalar(const Value* const value);
|
|
|
|
TORCH_CUDA_API c10::optional<float> getFloat(const Value* const value);
|
|
|
|
TORCH_CUDA_API c10::optional<int> getInt(const Value* const value);
|
|
|
|
// Returns the scalar as a float, regardless of its scalar type
|
|
// TODO: remove me
|
|
TORCH_CUDA_API float getAsFloat(const ::torch::jit::Value* const value);
|
|
|
|
/*
|
|
* Functions for comparing complete tensors
|
|
*/
|
|
|
|
TORCH_CUDA_API bool haveSameDevice(
|
|
const std::shared_ptr<c10::TensorType>& lhs,
|
|
const std::shared_ptr<c10::TensorType>& rhs);
|
|
|
|
TORCH_CUDA_API bool haveSameScalarType(
|
|
const std::shared_ptr<c10::TensorType>& lhs,
|
|
const std::shared_ptr<c10::TensorType>& rhs);
|
|
|
|
TORCH_CUDA_API bool haveSameSizes(
|
|
const std::shared_ptr<c10::TensorType>& lhs,
|
|
const std::shared_ptr<c10::TensorType>& rhs);
|
|
|
|
TORCH_CUDA_API bool haveSameStrides(
|
|
const std::shared_ptr<c10::TensorType>& lhs,
|
|
const std::shared_ptr<c10::TensorType>& rhs);
|
|
|
|
TORCH_CUDA_API bool haveSameShape(
|
|
const std::shared_ptr<c10::TensorType>& lhs,
|
|
const std::shared_ptr<c10::TensorType>& rhs);
|
|
|
|
} // namespace fuser
|
|
} // namespace jit
|
|
} // namespace torch
|