MCPcopy Create free account
hub / github.com/Rust-GPU/rust-cuda / binary_tensor_op

Method binary_tensor_op

crates/cudnn/src/context.rs:153–198  ·  view source on GitHub ↗

This function computes a binary element-wise tensor core operations according to the following equation: C = OP( alpha * A , beta * B ) + gamma * C given the tensors A, B and C, and the scaling parameters alpha, beta and gamma. Each dimension of the input tensor A must match the corresponding dimension of the destination tensor C, and each dimension of the input tensor B must match the correspo

(
        &self,
        op_desc: &OpTensorDescriptor<CompT, Op>,
        alpha: CompT,
        a: &Tensor<T1, F1, impl GpuBuffer<T1>, D>,
        beta: CompT,
        b: &Tensor<T2, F2, impl GpuBuffe

Source from the content-addressed store, hash-verified

151 /// to dimension five (5) are supported. This routine does not support tensor formats beyond
152 /// these dimensions.
153 pub fn binary_tensor_op<CompT, Op, T1, F1, T2, F2, T3, F3, const D: usize>(
154 &self,
155 op_desc: &OpTensorDescriptor<CompT, Op>,
156 alpha: CompT,
157 a: &Tensor<T1, F1, impl GpuBuffer<T1>, D>,
158 beta: CompT,
159 b: &Tensor<T2, F2, impl GpuBuffer<T2>, D>,
160 gamma: CompT,
161 c: &mut Tensor<T3, F3, impl GpuBuffer<T3>, D>,
162 ) -> Result<(), CudnnError>
163 where
164 CompT: DataType + SupportedOp<T1, T2, T3>,
165 Op: OpTensorOp + BinaryOp,
166 T1: DataType,
167 F1: TensorFormat + SupportedType<T1>,
168 T2: DataType,
169 F2: TensorFormat + SupportedType<T2>,
170 T3: DataType,
171 F3: TensorFormat + SupportedType<T3>,
172 {
173 let a_data = a.data().as_device_ptr().as_raw();
174 let a_desc = a.descriptor();
175
176 let b_data = b.data().as_device_ptr().as_raw();
177 let b_desc = b.descriptor();
178
179 let c_data = c.data().as_device_ptr().as_raw();
180 let c_desc = c.descriptor();
181
182 unsafe {
183 sys::cudnnOpTensor(
184 self.raw,
185 op_desc.raw,
186 &alpha as *const CompT as *const std::ffi::c_void,
187 a_desc.raw,
188 a_data as *const std::ffi::c_void,
189 &beta as *const CompT as *const std::ffi::c_void,
190 b_desc.raw,
191 b_data as *const std::ffi::c_void,
192 &gamma as *const CompT as *const std::ffi::c_void,
193 c_desc.raw,
194 c_data as *mut std::ffi::c_void,
195 )
196 .into_result()
197 }
198 }
199
200 /// This function computes an unary element wise tensor core operation according to the
201 /// following equation:

Callers

nothing calls this directly

Calls 5

as_rawMethod · 0.45
as_device_ptrMethod · 0.45
dataMethod · 0.45
descriptorMethod · 0.45
into_resultMethod · 0.45

Tested by

no test coverage detected