Understanding NMODL AST Transformations for LLVM branch

[pramodk]

What are AST transformations? See #474 #475 #476

There will be some additional transformation needed but before jumping into this we should understand current implementations. Here is what we should do:

• take nmodl and switch to llvm branch
• see how to build the llvm branch : https://github.com/BlueBrain/nmodl/blob/llvm/INSTALL.md#build-project . You can also see Jenkins job for BB5 : https://bbpcode.epfl.ch/ci/job/hpc.nmodl/configure or
• then see CodegenLLVMHelperVisitor class & implementation

How to play with current implementation?

• Create a sample mod file

$ cat hh.mod
TITLE hh.mod   squid sodium, potassium, and leak channels

UNITS {
    (mA) = (milliamp)
    (mV) = (millivolt)
    (S) = (siemens)
}

NEURON {
    SUFFIX hh
    USEION na READ ena WRITE ina
    USEION k READ ek WRITE ik
    NONSPECIFIC_CURRENT il
    RANGE gnabar, gkbar, gl, el, gna, gk
    RANGE minf, hinf, ninf, mtau, htau, ntau
    THREADSAFE : assigned GLOBALs will be per thread
}

PARAMETER {
    gnabar = .12 (S/cm2)    <0,1e9>
    gkbar = .036 (S/cm2)    <0,1e9>
    gl = .0003 (S/cm2)    <0,1e9>
    el = -54.3 (mV)
}

STATE {
    m h n
}

ASSIGNED {
    v (mV)
    celsius (degC)
    ena (mV)
    ek (mV)
    gna (S/cm2)
    gk (S/cm2)
    ina (mA/cm2)
    ik (mA/cm2)
    il (mA/cm2)
    minf hinf ninf
    mtau (ms) htau (ms) ntau (ms)
}

BREAKPOINT {
    SOLVE states METHOD cnexp
    gna = gnabar*m*m*m*h
    ina = gna*(v - ena)
    gk = gkbar*n*n*n*n
    ik = gk*(v - ek)
    il = gl*(v - el)
}

DERIVATIVE states {
     m' =  (minf-m)/mtau
     h' = (hinf-h)/htau
     n' = (ninf-n)/ntau
}

• In order to print transformation of AST, I did following change locally:

diff --git a/src/codegen/llvm/codegen_llvm_helper_visitor.cpp b/src/codegen/llvm/codegen_llvm_helper_visitor.cpp
index fc8fda3..4d0f17f 100644
--- a/src/codegen/llvm/codegen_llvm_helper_visitor.cpp
+++ b/src/codegen/llvm/codegen_llvm_helper_visitor.cpp
@@ -428,6 +428,8 @@ void CodegenLLVMHelperVisitor::visit_nrn_state_block(ast::NrnStateBlock& node) {
     /// statements for new function to be generated
     ast::StatementVector function_statements;

+    std::cout << "BEFORE :: \n" << nmodl::to_nmodl(node) << "\n";
+
     /// create variable definition for loop index and insert at the beginning
     std::string loop_index_var = "id";
     std::vector<std::string> int_variables{"id"};
@@ -527,7 +529,7 @@ void CodegenLLVMHelperVisitor::visit_nrn_state_block(ast::NrnStateBlock& node) {
         std::make_shared<ast::CodegenFunction>(return_type, name, code_arguments, function_block);
     codegen_functions.push_back(function);

-    std::cout << nmodl::to_nmodl(function);
+    std::cout << "AFTER :: \n" << nmodl::to_nmodl(function) << "\n";
 }

• Build the project and now run nmodl binary on hh.mod file provided above:

$ ./bin/nmodl hh.mod llvm --ir
[NMODL] [info] :: Processing hh.mod
[NMODL] [info] :: Running symtab visitor
[NMODL] [info] :: Running CVode to cnexp visitor
[NMODL] [info] :: Running LOCAL to ASSIGNED visitor
[NMODL] [info] :: Running code compatibility checker
[NMODL] [info] :: Running verbatim rename visitor
[NMODL] [info] :: Running KINETIC block visitor
[NMODL] [info] :: Running STEADYSTATE visitor
[NMODL] [info] :: Parsing Units
[NMODL] [info] :: Running cnexp visitor
[NMODL] [info] :: Running C backend code generator
[NMODL] [info] :: Running LLVM backend code generator
[NMODL] [info] :: Running CodegenLLVMHelperVisitor

BEFORE ::
NRN_STATE SOLVE states METHOD cnexp{
    m = m+(1.0-exp(dt*((((-1.0)))/mtau)))*(-(((minf))/mtau)/((((-1.0)))/mtau)-m)
    h = h+(1.0-exp(dt*((((-1.0)))/htau)))*(-(((hinf))/htau)/((((-1.0)))/htau)-h)
    n = n+(1.0-exp(dt*((((-1.0)))/ntau)))*(-(((ninf))/ntau)/((((-1.0)))/ntau)-n)
}

AFTER ::
VOID nrn_state_hh(){
    INTEGER id
    for(id = 0; id<node_count; id = id+1) {
        INTEGER node_id, ena_id, ek_id
        DOUBLE v
        node_id = node_index[id]
        ena_id = ion_ena_index[id]
        ek_id = ion_ek_index[id]
        v = voltage[node_id]
        ena[id] = ion_ena[ena_id]
        ek[id] = ion_ek[ek_id]
        m[id] = m[id]+(1.0-exp(dt*((((-1.0)))/mtau[id])))*(-(((minf[id]))/mtau[id])/((((-1.0)))/mtau[id])-m[id])
        h[id] = h[id]+(1.0-exp(dt*((((-1.0)))/htau[id])))*(-(((hinf[id]))/htau[id])/((((-1.0)))/htau[id])-h[id])
        n[id] = n[id]+(1.0-exp(dt*((((-1.0)))/ntau[id])))*(-(((ninf[id]))/ntau[id])/((((-1.0)))/ntau[id])-n[id])
    }
}
; ModuleID = 'hh'
.....ignore this IR output.....

See BEFORE and AFTER output.

• What we are interested in is how BEFORE block gets translated into AFTER block. This transformation is done in https://github.com/BlueBrain/nmodl/blob/6c72107969c3a3196ee86cc3e2a082b6cad48c09/src/codegen/llvm/codegen_llvm_helper_visitor.cpp#L427

• What we have to see is:

  • how ast::NrnStateBlock is cloned
  • new expression creations
  • new statements creations
  • new block creation
  • new for loop construction

cc: @alkino

[pramodk]

Important : How to debug AST ?

nmodl::to_nmodl(node) doesn't show full AST structure but just text. Note that to_nmodl may not show certain nodes. In case of doubt, see <build>/src/visitors/nmodl_visitor.cpp which transforms AST to NMODL i.e. txt form.

Better way to inspect AST is using nmodl::to_json(const nmodl::ast::Ast& node); function. Using following change:

+++ b/src/codegen/llvm/codegen_llvm_helper_visitor.cpp
@@ -428,6 +428,8 @@ void CodegenLLVMHelperVisitor::visit_nrn_state_block(ast::NrnStateBlock& node) {
     /// statements for new function to be generated
     ast::StatementVector function_statements;

.......

-    std::cout << nmodl::to_nmodl(function);
+    std::cout << "AFTER :: \n" << nmodl::to_nmodl(function) << "\n";
+    std::cout << "AFTER (json) :: \n" << nmodl::to_json(*function) << "\n";
 }

and re-running nmodl on hh.mod will show:

$ ./bin/nmodl hh.mod llvm --ir
.....

AFTER ::
VOID nrn_state_hh(){
    INTEGER id
    for(id = 0; id<node_count; id = id+1) {
        INTEGER node_id, ena_id, ek_id
        DOUBLE v
        node_id = node_index[id]
        ena_id = ion_ena_index[id]
        ek_id = ion_ek_index[id]
        v = voltage[node_id]
        ena[id] = ion_ena[ena_id]
        ek[id] = ion_ek[ek_id]
        m[id] = m[id]+(1.0-exp(dt*((((-1.0)))/mtau[id])))*(-(((minf[id]))/mtau[id])/((((-1.0)))/mtau[id])-m[id])
        h[id] = h[id]+(1.0-exp(dt*((((-1.0)))/htau[id])))*(-(((hinf[id]))/htau[id])/((((-1.0)))/htau[id])-h[id])
        n[id] = n[id]+(1.0-exp(dt*((((-1.0)))/ntau[id])))*(-(((ninf[id]))/ntau[id])/((((-1.0)))/ntau[id])-n[id])
    }
}

AFTER (json) ::
{
  "CodegenFunction": [
    {
      "CodegenVarType": [
        {
          "name": "VOID"
        }
      ]
    },
    {
      "Name": [
        {
          "String": [
            {
              "name": "nrn_state_hh"
            }
          ]
        }
      ]
    },
    {
      "StatementBlock": [
        {
          "CodegenVarListStatement": [
            {
              "CodegenVarType": [
                {
                  "name": "INTEGER"
                }
              ]
            },
            {
              "CodegenVar": [
                {
                  "Name": [
                    {
                      "String": [
                        {
                          "name": "id"
                        }
....

The JSON form printed above represent every construct in the AST. This is quite helpful to inspect whole AST or debug issues.

[pramodk]

@alkino : once this is done / you get an idea of overall visitor, please close this.