softmax namespace, shape & type inference

daphne-eu · Jul 1, 2024 · 1ae98c5 · 1ae98c5
1 parent 555269d
commit 1ae98c5
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 deletions.
diff --git a/src/ir/daphneir/DaphneOps.td b/src/ir/daphneir/DaphneOps.td
@@ -900,7 +900,7 @@ def Daphne_AvgPoolForwardOp : Daphne_PoolForwardOp<"NN_Avg_Forward", [ DeclareOp
 // ----------------------------------------------------------------------------
 // Softmax
 // ----------------------------------------------------------------------------
-def Daphne_SoftmaxForwardOp : Daphne_Op<"softmaxForward", [ CUDASupport ]> {
+def Daphne_SoftmaxForwardOp : Daphne_Op<"NN_Softmax_Forward", [ TypeFromFirstArg, ShapeFromArg, CUDASupport ]> {
     let arguments = (ins  MatrixOf<[FloatScalar]>:$input);
 
     let results = (outs MatrixOf<[FloatScalar]>:$data);

diff --git a/src/runtime/local/kernels/CUDA/Softmax.cpp b/src/runtime/local/kernels/CUDA/Softmax.cpp
@@ -17,7 +17,7 @@
 #include "Softmax.h"
 #include <runtime/local/datastructures/AllocationDescriptorCUDA.h>
 
-namespace CUDA::Softmax {
+namespace CUDA::NN::Softmax {
 
     template<typename DTRes, typename DTArg>
     void Forward<DTRes, DTArg>::apply(DTRes *&res, const DTArg *data, DCTX(dctx)) {

diff --git a/src/runtime/local/kernels/CUDA/Softmax.h b/src/runtime/local/kernels/CUDA/Softmax.h
@@ -28,7 +28,7 @@
 #include <cstddef>
 #include <cstdint>
 
-namespace CUDA::Softmax {
+namespace CUDA::NN::Softmax {
     template<typename DTRes, typename DTArg>
     struct Forward {
         static void apply(DTRes *&res, const DTArg *data, DCTX(dctx));

diff --git a/src/runtime/local/kernels/kernels.json b/src/runtime/local/kernels/kernels.json
@@ -3915,7 +3915,7 @@
     {
         "kernelTemplate": {
             "header": "Softmax.h",
-            "opName": "Softmax::Forward",
+            "opName": "NN::Softmax::Forward",
             "returnType": "void",
             "opCodeAsTemplateParam": 1,
             "templateParams": [