// Copyright (c) 2024 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. #include "paddle/phi/backends/gpu/gpu_context.h" #include "paddle/phi/core/kernel_registry.h" #include "paddle/phi/core/visit_type.h" #include "paddle/phi/kernels/funcs/sparse/convolution.h" #include "paddle/phi/kernels/funcs/transpose_function.cuh" #include "paddle/phi/kernels/sparse/gpu/conv_kernel_impl.cuh" #include "paddle/phi/kernels/sparse/gpu/sparse_conv_hashmap.cuh" #include "glog/logging.h" namespace phi { namespace sparse { template void Conv3dImplicitGemmGPUKernel(const GPUContext& dev_ctx, const SparseCooTensor& x, const DenseTensor& kernel, const std::vector& paddings, const std::vector& dilations, const std::vector& strides, const int groups, const bool subm, const std::string& key, SparseCooTensor* out) { // Currently, only support x.layout is NDHWC, subm = true, stride = 1, groups // = 1, dilations = 1 PADDLE_ENFORCE_EQ( subm, true, common::errors::InvalidArgument("The subm must be true, but received %s.", subm ? "true" : "false")); PADDLE_ENFORCE_EQ(groups, 1, common::errors::InvalidArgument( "The group must be 1, but received %d.", groups)); const auto& x_dims = x.dims(); const auto& kernel_dims = kernel.dims(); const bool is2D = x_dims.size() == 4 ? true : false; if (is2D) { PADDLE_ENFORCE_EQ( (kernel_dims.size() == 4), true, common::errors::InvalidArgument( "For 2D case, the size of kernel_dims must be 4, but received %d.", kernel_dims.size())); PADDLE_ENFORCE_EQ( (strides.size() == 2 && strides[0] == 1 && strides[1] == 1), true, common::errors::InvalidArgument( "The strides must be 1, but received %d, %d.", strides[0], strides[1])); PADDLE_ENFORCE_EQ( (dilations.size() == 2 && dilations[0] == 1 && dilations[1] == 1), true, common::errors::InvalidArgument( "The dilations must be 1, but received %d, %d.", dilations[0], dilations[1])); } else { PADDLE_ENFORCE_EQ( (kernel_dims.size() == 5), true, common::errors::InvalidArgument( "For 3D case, the size of kernel_dims must be 5, but received %d.", kernel_dims.size())); PADDLE_ENFORCE_EQ((strides.size() == 3 && strides[0] == 1 && strides[1] == 1 && strides[2] == 1), true, common::errors::InvalidArgument( "The strides must be 1, but received %d, %d, %d.", strides[0], strides[1], strides[2])); PADDLE_ENFORCE_EQ((dilations.size() == 3 && dilations[0] == 1 && dilations[1] == 1 && dilations[2] == 1), true, common::errors::InvalidArgument( "The dilations must be 1, but received %d, %d, %d.", dilations[0], dilations[1], dilations[2])); } int kernel_volume = is2D ? kernel_dims[0] * kernel_dims[1] : kernel_dims[0] * kernel_dims[1] * kernel_dims[2]; int in_channels = is2D ? kernel_dims[2] : kernel_dims[3]; int out_channels = is2D ? kernel_dims[3] : kernel_dims[4]; int rank = is2D ? 4 : 5; std::vector out_dims_vec(rank, 1); DDim out_dims = make_ddim(out_dims_vec); std::vector kernel_sizes(kernel_dims.size()); for (int i = 0; i < kernel_dims.size(); i++) { kernel_sizes[i] = kernel_dims[i]; } std::vector subm_paddings(paddings), subm_strides(strides); if (subm) { // the out shape of subm_conv is same as input shape // reset the padding=kernel_size/2 and strides=1 funcs::sparse::ResetSubmKernelSizeAndStrides( kernel.dims(), &subm_paddings, &subm_strides); } funcs::sparse::GetOutShape( x_dims, kernel_sizes, subm_paddings, dilations, subm_strides, &out_dims); // Set the output tensor if (subm) { DenseTensor out_indices = EmptyLike(dev_ctx, x.indices()); int tmpidx = is2D ? 3 : 4; DenseTensor out_values = Empty(dev_ctx, {x.nnz(), kernel_sizes[tmpidx]}); phi::Copy(dev_ctx, x.indices(), dev_ctx.GetPlace(), false, &out_indices); out->SetMember(out_indices, out_values, out_dims, false); } else { PADDLE_THROW(common::errors::Unimplemented( "The subm must be true, but received %s.", subm ? "true" : "false")); } build_sparse_conv_kmap( dev_ctx, x, key, kernel_sizes, strides, kernel_volume, is2D, out); auto* out_kmap_cache_ptr = out->GetKmapCache(key); DenseTensor kernel_transpose = EmptyLike(dev_ctx, kernel); std::vector perm; if (is2D) { perm = {1, 0, 2, 3}; } else { perm = {2, 1, 0, 3, 4}; } funcs::TransposeGPUKernelDriver(dev_ctx, kernel, perm, &kernel_transpose); #ifdef PADDLE_WITH_CUDA conv_forward_implicit_gemm_cuda(dev_ctx, x.values(), kernel_transpose, *(out_kmap_cache_ptr->out_in_map), out->nnz(), out_channels, *(out->mutable_values())); #else PADDLE_THROW(common::errors::Unimplemented( "conv_forward_implicit_gemm_cuda is only supported on CUDA.")); #endif } /** * x: the input SparseCooTensor, shape is (N, D, H, W, C) * kernel: the weight data, shape is (D, H, W, C, OC) * out: the output SparseCooTensor, shape is (N, D, H, W, OC) * rulebook: return rulebook if key is not valid else return nullptr * counter: return counter if key is not valid else return nullptr **/ template void Conv3dImplicitGemmKernel(const Context& dev_ctx, const SparseCooTensor& x, const DenseTensor& kernel, const std::vector& paddings, const std::vector& dilations, const std::vector& strides, const int groups, const bool subm, const std::string& key, SparseCooTensor* out) { #ifdef PADDLE_WITH_CUDA PD_VISIT_BASE_INTEGRAL_TYPES( x.indices().dtype(), "Conv3dImplicitGemmGPUKernel's indices", ([&] { // Conv3dImplicitGemmGPUKernel(dev_ctx, Conv3dImplicitGemmGPUKernel(dev_ctx, x, kernel, paddings, dilations, strides, groups, subm, key, out); })); #else PADDLE_THROW(common::errors::Unimplemented( "Conv3dImplicitGemmKernel is only supported on CUDA.")); #endif } } // namespace sparse } // namespace phi PD_REGISTER_KERNEL(conv3d_implicit_gemm, GPU, ALL_LAYOUT, phi::sparse::Conv3dImplicitGemmKernel, float, phi::float16) { kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO); kernel->OutputAt(0).SetDataType(phi::DataType::UNDEFINED); }