Fix type bug in conv2d/gemm with broadcast (#796)

add ElementVector

---------

Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
Shuai Shao
2023-02-09 20:53:25 -05:00
committed by GitHub
co-authored by Haicheng Wu
parent 2e10404d26
commit ce8597dc14
7 changed files with 21 additions and 11 deletions
@@ -120,6 +120,7 @@ public:
using EpilogueOutputOp = typename Conv2d::EpilogueOutputOp;
using ElementZ = typename EpilogueOutputOp::ElementZ;
using ElementT = typename EpilogueOutputOp::ElementT;
using ElementVector = typename EpilogueOutputOp::ElementVector;
static cutlass::conv::Operator const kConvolutionalOperator = Conv2d::kConvolutionalOperator;
static const bool kAddBroadcastFirst = AddBroadcastFirst;
@@ -142,7 +143,7 @@ public:
cutlass::HostTensor<ElementT, LayoutC> tensor_T_computed;
cutlass::HostTensor<ElementT, LayoutC> tensor_T_reference;
cutlass::HostTensor<ElementAccumulator, LayoutC> tensor_Y_reference;
cutlass::HostTensor<ElementC, LayoutC> tensor_Broadcast; // Input Broadcast
cutlass::HostTensor<ElementVector, LayoutC> tensor_Broadcast; // Input Broadcast
public:
@@ -105,7 +105,8 @@ struct TestbedGemmWithBroadcast {
using OutputOp = typename Gemm::GemmKernel::Epilogue::OutputOp;
using ElementC = typename Gemm::ElementC;
using ElementAccumulator = typename Gemm::ElementAccumulator;
using ElementCOmpute = typename OutputOp::ElementCompute;
using ElementCompute = typename OutputOp::ElementCompute;
using ElementVector = typename OutputOp::ElementVector;
using ElementZ = typename OutputOp::ElementZ;
using ElementT = typename OutputOp::ElementT;
@@ -118,7 +119,7 @@ struct TestbedGemmWithBroadcast {
cutlass::HostTensor<typename Gemm::ElementA, typename Gemm::LayoutA> tensor_A; // Input A
cutlass::HostTensor<typename Gemm::ElementB, typename Gemm::LayoutB> tensor_B; // Input B
cutlass::HostTensor<ElementC, typename Gemm::LayoutC> tensor_C; // Input C
cutlass::HostTensor<ElementC, typename Gemm::LayoutC> tensor_Broadcast; // Input Broadcast
cutlass::HostTensor<ElementVector, typename Gemm::LayoutC> tensor_Broadcast; // Input Broadcast
cutlass::HostTensor<ElementZ, typename Gemm::LayoutC> tensor_Z;
cutlass::HostTensor<ElementT, typename Gemm::LayoutC> tensor_T;