ZLUDA v3.9.1 (#84)

* Handle ptx kernel parameter attribute .ptr.

* Add cublasStrmm_v2, cublasSgeam.

* Handle MIOpen internal error.

* Implement cuStreamGetPriority.

* Add cudnnBatchNormalizationForwardInference.

* Fix nvrtc.

Author: lshqqytiger

ptx/src/ast.rs

@@ -144,6 +144,10 @@ pub struct Function<'a, ID, S> {
     pub body: Option<Vec<S>>,
 }
 
+pub enum KernelParameterAttribute {
+    Pointer(StateSpace),
+}
+
 #[derive(PartialEq, Eq, Clone, Hash)]
 pub enum Type {
     // .param.b32 foo;

ptx/src/ptx.lalrpop

@@ -124,6 +124,7 @@ match {
     ".popc",
     ".pragma",
     ".pred",
+    ".ptr",
     ".r",
     ".red",
     ".reg",
@@ -566,6 +567,13 @@ KernelArguments: Vec<ast::VariableDeclaration<&'input str>> = {
     "(" <args:Comma<VariableDeclarationEntry>> ")" => args
 };
 
+// https://docs.nvidia.com/cuda/parallel-thread-execution/#kernel-function-parameter-attributes
+KernelParameterAttribute: ast::KernelParameterAttribute = {
+    ".ptr" <state_space:StateSpaceSpecifier?> => {
+        ast::KernelParameterAttribute::Pointer(state_space.unwrap_or(ast::StateSpace::Generic))
+    }
+};
+
 FnArguments: Vec<ast::VariableDeclaration<&'input str>> = {
     "(" <args:Comma<VariableDeclarationFunc>> ")" => args
 };
@@ -662,21 +670,21 @@ InitializerNoAdd: ast::Initializer<&'input str> = {
 }
 
 VariableDeclarationFunc: ast::VariableDeclaration<&'input str> = {
-    <var:VariableDeclarationBase> =>
+    <var:VariableDeclaration> =>
         validate_variable_declaration_func(var, errors)
 }
 
 VariableDeclarationEntry: ast::VariableDeclaration<&'input str> = {
-    <var:VariableDeclarationBase> =>
+    <var:VariableDeclaration> =>
         validate_variable_declaration_entry(var, errors)
 }
 
 VariableDeclarationProto: ast::VariableDeclaration<&'input str> = {
-    <var:VariableDeclarationBase> =>
+    <var:VariableDeclaration> =>
         validate_variable_declaration_proto(var, errors)
 }
 
-VariableDeclarationBase: ast::VariableDeclaration<&'input str> = {
+VariableDeclaration: ast::VariableDeclaration<&'input str> = {
     <variable:VariableDeclarationBegin> <name:ExtendedIDOrBlank> <dims:ArrayDimensions?> => {
         let mut variable = variable.clone();
         variable.name = name;
@@ -710,16 +718,54 @@ MultiVariableDefinition: Vec<ast::MultiVariableDefinition<&'input str>> = {
 }
 
 VariableDeclarationBegin: ast::VariableDeclaration<&'input str> = {
-    <state_space:StateSpaceSpecifier> <align:Align?> <type_:AnyType> => {
-        ast::VariableDeclaration {
+    <state_space:StateSpaceSpecifier> <align:Align> <v_len:VectorPrefix?> <type_:ScalarType> => {
+        let mut variable = ast::VariableDeclaration {
+            align: Some(align),
+            type_: ast::Type::Scalar(type_),
+            state_space,
+            name: ""
+        };
+        if let Some(v_len) = v_len {
+            variable.type_ = ast::Type::Vector(type_, v_len)
+        }
+        variable
+    },
+    <state_space:StateSpaceSpecifier> <v_len:VectorPrefix?> <type_:ScalarType> <attr:KernelParameterAttribute?> <align:Align?> => {
+        let mut variable = ast::VariableDeclaration {
             align,
-            type_,
+            type_: ast::Type::Scalar(type_),
+            state_space,
+            name: ""
+        };
+        if let Some(v_len) = v_len {
+            variable.type_ = ast::Type::Vector(type_, v_len)
+        }
+        if let Some(attr) = attr {
+            variable.type_ = match attr {
+                ast::KernelParameterAttribute::Pointer(state_space) => {
+                    ast::Type::Pointer(type_, state_space)
+                }
+            };
+        }
+        variable
+    },
+    <state_space:StateSpaceSpecifier> <var:VariableDeclarationBeginNonScalarTypeWithAlign> => {
+        ast::VariableDeclaration {
+            align: var.1,
+            type_: var.0,
             state_space,
             name: ""
         }
     }
 }
 
+VariableDeclarationBeginNonScalarTypeWithAlign: (ast::Type, Option<u32>) = {
+    <align:Align> ".texref" => (ast::Type::Texref, Some(align)),
+    ".texref" <align:Align?> => (ast::Type::Texref, align),
+    <align:Align> ".surfref" => (ast::Type::Surfref, Some(align)),
+    ".surfref" <align:Align?> => (ast::Type::Surfref, align)
+}
+
 VariableDefinitionOnce: (&'input str, Option<Either<u32, Vec<u32>>>, Option<ast::Initializer<&'input str>>) = {
     <name:ExtendedIDOrBlank> <suffix:VariableDefinitionSuffix?> <init:VariableDefinitionInitializer?> => (name, suffix, init)
 }
@@ -734,13 +780,6 @@ VariableDefinitionInitializer: ast::Initializer<&'input str> = {
     "=" <init:Initializer> => init
 }
 
-AnyType: ast::Type = {
-    ".texref" => ast::Type::Texref,
-    ".surfref" => ast::Type::Surfref,
-    <v_len:VectorPrefix> <type_:ScalarType> => ast::Type::Vector(type_, v_len),
-    <type_:ScalarType> => ast::Type::Scalar(type_),
-}
-
 #[inline]
 SizedScalarType: ast::ScalarType = {
     ".b8" => ast::ScalarType::B8,
@@ -2780,58 +2819,6 @@ AnyBitType: ast::ScalarType = {
     ".b64" => ast::ScalarType::B64,
 };
 
-VariableScalarUnitialized<T>: (Option<u32>, T, &'input str) = {
-    <align:Align?> <type_:T> <name:ExtendedID> => {
-        (align, type_, name)
-    }
-}
-
-VariableScalar<T>: (Option<u32>, T, &'input str, Vec<u8>) = {
-    <align:Align?> <type_:T> <name:ExtendedID> <init:VariableInitalizer?> => {
-        let initializer = init.map(ast::ImmediateValue::to_bytes).unwrap_or(Vec::new());
-        (align, type_, name, initializer)
-    }
-}
-
-VariableInitalizer: ast::ImmediateValue = {
-    "=" <v:ImmediateValue> => v
-}
-
-VariableVector<T>: (Option<u32>, u8, T, &'input str) = {
-    <align:Align?> <v_len:VectorPrefix> <type_:T> <name:ExtendedID> => {
-        (align, v_len, type_, name)
-    }
-}
-
-// empty dimensions [0] means it's a pointer
-VariableArrayOrPointer<T>: (Option<u32>, T, &'input str, ast::ArrayOrPointer) = {
-    <align:Align?> <typ:SizedScalarType> <name:ExtendedID> <dims:ArrayDimensions> <init:ArrayInitializer?> => {
-        let mut dims = dims;
-        let array_init = match init {
-            Some(init) => {
-                let init_vec = match init.to_vec(typ, &mut dims) {
-                    Err(error) => {
-                        errors.push(ParseError::User { error });
-                        Vec::new()
-                    }
-                    Ok(x) => x
-                };
-                ast::ArrayOrPointer::Array { dimensions: dims, init: init_vec }
-            }
-            None => {
-                if dims.len() > 1 && dims.contains(&0) {
-                    errors.push(ParseError::User { error: ast::PtxError::ZeroDimensionArray });
-                }
-                match &*dims {
-                    [0] => ast::ArrayOrPointer::Pointer,
-                    _ => ast::ArrayOrPointer::Array { dimensions: dims, init: Vec::new() }
-                }
-            }
-        };
-        (align, typ, name, array_init)
-    }
-}
-
 // [0] and [] are treated the same
 ArrayDimensions: Vec<u32> = {
     ArrayEmptyDimension => vec![0u32],

zluda/src/cuda.rs

@@ -144,6 +144,7 @@ cuda_function_declarations!(
         cuStreamGetCtx,
         cuStreamGetCtx_ptsz,
         cuStreamGetFlags,
+        cuStreamGetPriority,
         cuStreamIsCapturing,
         cuStreamQuery,
         cuStreamSynchronize,
@@ -1087,6 +1088,13 @@ mod definitions {
         stream::get_flags(hStream, flags)
     }
 
+    pub(crate) unsafe fn cuStreamGetPriority(
+        hStream: *mut stream::Stream,
+        priority: *mut ::std::os::raw::c_int,
+    ) -> Result<(), CUresult> {
+        stream::get_priority(hStream, priority)
+    }
+
     pub(crate) unsafe fn cuStreamIsCapturing(
         hStream: *mut stream::Stream,
         captureStatus: *mut hipStreamCaptureStatus,
@@ -1662,9 +1670,7 @@ mod definitions {
         array::mipmapped_create(pHandle, pMipmappedArrayDesc, numMipmapLevels)
     }
 
-    pub(crate) unsafe fn cuMipmappedArrayDestroy(
-        hMipmappedArray: CUmipmappedArray,
-    ) -> hipError_t {
+    pub(crate) unsafe fn cuMipmappedArrayDestroy(hMipmappedArray: CUmipmappedArray) -> hipError_t {
         array::mipmapped_destroy(hMipmappedArray)
     }
 

zluda/src/impl/stream.rs

@@ -181,6 +181,12 @@ pub(crate) unsafe fn get_flags(stream: *mut Stream, flags: *mut u32) -> Result<(
     Ok(())
 }
 
+pub(crate) unsafe fn get_priority(stream: *mut Stream, priority: *mut i32) -> Result<(), CUresult> {
+    let hip_stream = as_hip_stream(stream)?;
+    hip_call_cuda! { hipStreamGetPriority(hip_stream, priority) };
+    Ok(())
+}
+
 pub(crate) unsafe fn is_capturing(
     stream: *mut Stream,
     capture_status: *mut hipStreamCaptureStatus,

zluda_blas/src/cublas.rs

@@ -3828,7 +3828,7 @@ pub extern "system" fn cublasZtrsm_v2(
 }
 
 #[no_mangle]
-pub extern "system" fn cublasStrmm_v2(
+pub unsafe extern "system" fn cublasStrmm_v2(
     handle: cublasHandle_t,
     side: cublasSideMode_t,
     uplo: cublasFillMode_t,
@@ -3844,7 +3844,9 @@ pub extern "system" fn cublasStrmm_v2(
     C: *mut f32,
     ldc: ::std::os::raw::c_int,
 ) -> cublasStatus_t {
-    crate::unsupported()
+    crate::strmm(
+        handle, side, uplo, trans, diag, m, n, alpha, A, lda, B, ldb, C, ldc,
+    )
 }
 
 #[no_mangle]
@@ -4292,7 +4294,7 @@ pub extern "system" fn cublasHgemmStridedBatched(
 }
 
 #[no_mangle]
-pub extern "system" fn cublasSgeam(
+pub unsafe extern "system" fn cublasSgeam(
     handle: cublasHandle_t,
     transa: cublasOperation_t,
     transb: cublasOperation_t,
@@ -4307,7 +4309,9 @@ pub extern "system" fn cublasSgeam(
     C: *mut f32,
     ldc: ::std::os::raw::c_int,
 ) -> cublasStatus_t {
-    crate::unsupported()
+    crate::sgeam(
+        handle, transa, transb, m, n, alpha, A, lda, beta, B, ldb, C, ldc,
+    )
 }
 
 #[no_mangle]

zluda_blas/src/lib.rs

@@ -489,6 +489,40 @@ fn to_compute_type(compute_type: cublasComputeType_t) -> rocblas_datatype {
     }
 }
 
+unsafe fn sgeam(
+    handle: *mut cublasContext,
+    transa: cublasOperation_t,
+    transb: cublasOperation_t,
+    m: i32,
+    n: i32,
+    alpha: *const f32,
+    a: *const f32,
+    lda: i32,
+    beta: *const f32,
+    b: *const f32,
+    ldb: i32,
+    c: *mut f32,
+    ldc: i32,
+) -> cublasStatus_t {
+    let transa = op_from_cuda(transa);
+    let transb = op_from_cuda(transb);
+    to_cuda(rocblas_sgeam(
+        handle.cast(),
+        transa,
+        transb,
+        m,
+        n,
+        alpha,
+        a,
+        lda,
+        beta,
+        b,
+        ldb,
+        c,
+        ldc,
+    ))
+}
+
 unsafe fn zgemm_strided_batch(
     handle: *mut cublasContext,
     transa: cublasOperation_t,
@@ -1146,6 +1180,44 @@ unsafe fn dtrsm(
     ))
 }
 
+unsafe fn strmm(
+    handle: *mut cublasContext,
+    side: cublasSideMode_t,
+    uplo: cublasFillMode_t,
+    trans: cublasOperation_t,
+    diag: cublasDiagType_t,
+    m: i32,
+    n: i32,
+    alpha: *const f32,
+    a: *const f32,
+    lda: i32,
+    b: *const f32,
+    ldb: i32,
+    c: *mut f32,
+    ldc: i32,
+) -> cublasStatus_t {
+    let side = to_side(side);
+    let uplo = to_fill(uplo);
+    let trans = op_from_cuda(trans);
+    let diag = to_diag(diag);
+    to_cuda(rocblas_strmm(
+        handle.cast(),
+        side,
+        uplo,
+        trans,
+        diag,
+        m,
+        n,
+        alpha,
+        a,
+        lda,
+        b,
+        ldb,
+        c,
+        ldc,
+    ))
+}
+
 unsafe fn gemm_batched_ex(
     handle: cublasHandle_t,
     transa: cublasOperation_t,

zluda_dnn/src/cudnn.rs

@@ -2729,7 +2729,7 @@ impl cudnnBatchNormOps_t {
 pub struct cudnnBatchNormOps_t(pub ::std::os::raw::c_int);
 
 #[no_mangle]
-pub extern "system" fn cudnnBatchNormalizationForwardInference(
+pub unsafe extern "system" fn cudnnBatchNormalizationForwardInference(
     handle: cudnnHandle_t,
     mode: cudnnBatchNormMode_t,
     alpha: *const ::std::os::raw::c_void,
@@ -2745,7 +2745,22 @@ pub extern "system" fn cudnnBatchNormalizationForwardInference(
     estimatedVariance: *const ::std::os::raw::c_void,
     epsilon: f64,
 ) -> cudnnStatus_t {
-    crate::unsupported()
+    crate::batch_normalization_forward_inference(
+        handle,
+        mode,
+        alpha,
+        beta,
+        xDesc,
+        x,
+        yDesc,
+        y,
+        bnScaleBiasMeanVarDesc,
+        bnScale,
+        bnBias,
+        estimatedMean,
+        estimatedVariance,
+        epsilon,
+    )
 }
 impl cudnnNormMode_t {
     pub const CUDNN_NORM_PER_ACTIVATION: cudnnNormMode_t = cudnnNormMode_t(0);