polish fp16

Jittor · Mar 18, 2022 · ad57ec8 · ad57ec8
1 parent 5efb222
commit ad57ec8
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/python/jittor/extern/cuda/cublas/ops/cublas_batched_matmul_op.cc b/python/jittor/extern/cuda/cublas/ops/cublas_batched_matmul_op.cc
@@ -89,7 +89,7 @@ void CublasBatchedMatmulOp::jit_prepare(JK& jk) {
     jk << _CS("[T:") << a->dtype();
     jk << _CS("][Trans_a:") << (trans_a ? 'T' : 'N');
     jk << _CS("][Trans_b:") << (trans_b ? 'T' : 'N');
-    jk << _CS("][op:") << (a->dtype().dsize() == 4 ? 'S' : 'D');
+    jk << _CS("][op:") << (a->dtype().dsize() == 2? 'H' : (a->dtype().dsize() == 4 ? 'S' : 'D'));
     jk << ']';
 }
 

diff --git a/python/jittor/extern/cuda/cublas/ops/cublas_matmul_op.cc b/python/jittor/extern/cuda/cublas/ops/cublas_matmul_op.cc
@@ -50,7 +50,7 @@ void CublasMatmulOp::jit_prepare(JK& jk) {
     jk << _CS("[T:") << a->dtype();
     jk << _CS("][Trans_a:") << (trans_a ? 'T' : 'N');
     jk << _CS("][Trans_b:") << (trans_b ? 'T' : 'N');
-    jk << _CS("][op:") << (a->dtype().dsize() == 4 ? 'S' : 'D');
+    jk << _CS("][op:") << (a->dtype().dsize() == 2? 'H' : (a->dtype().dsize() == 4 ? 'S' : 'D'));
     jk << ']';
 }
 

diff --git a/python/jittor/other/code_softmax.py b/python/jittor/other/code_softmax.py
@@ -100,7 +100,7 @@ def softmax_v1(a, log=False):
     {for_loop} 
         #pragma unroll
         for (int j=0; j<{ILP}; j++)
-            v1 += {"vy[i][j];" if log else "vx[i][j]*vy[i][j];"}
+            v1 += {"float(vy[i][j]);" if log else "float(vx[i][j]*vy[i][j]);"}
 
     typedef cub::BlockReduce<float, {tnum}> BlockReduce;
     __shared__ typename BlockReduce::TempStorage temp_storage;
@@ -114,8 +114,8 @@ def softmax_v1(a, log=False):
         #pragma unroll
         for (int j=0; j<{ILP}; j++)
             vx[i][j] = {
-                "vy[i][j] - expf(vx[i][j]) * reduce_var;" if log 
-                else "vx[i][j] * (vy[i][j] - reduce_var);"
+                "vy[i][j] - in0_type(expf(vx[i][j]) * reduce_var);" if log 
+                else "vx[i][j] * (vy[i][j] - in0_type(reduce_var));"
             }
 
     {for_loop}