Support for Breakpoint block (nrn_cur) for code generation

[pramodk]

• similar to DERIVATIVE (nrn_state), handle BREAKPOINT (nrn_cur) blocks with AST level transformation
• refactor nrn_state & nrn_cur to use common code to avoid
• fix the llvm tests
• implement nrn_current calls similar to C codeine backend
• added tests

fixes #644

[pramodk]

@georgemitenkov : this is still WIP in terms of completing full/correct functionality of nrn_cur(). Also, I am going to refactor visit_breakpoint_block and visit_nrn_state_block in order to use common code to generate main loop, epilogue loop etc instead of duplication.

But all those changes will be inside CodegenLLVMHelperVisitor. So if you want, you can already take this branch and continue working on other part of LLVM code generation for nrn_cur function. For example, if you want to add test for nrn_cur, run it via JIT etc.

In my test example:

 NEURON {
     SUFFIX hh
     NONSPECIFIC_CURRENT il
     RANGE minf, mtau, gl, el
     USEION na WRITE nai
 }

 STATE {
     m
 }

 ASSIGNED {
     v (mV)
     minf
     mtau (ms)
     il
     el
     gl
 }

 BREAKPOINT {
     SOLVE states METHOD cnexp
     il = gl*(v - el)
 }

 DERIVATIVE states {
      m = minf + mtau + 0.1
      :printf("%lf %d %lf", m, ion_nai_id, nai)
 }

produces:

→ ./bin/nmodl ../test.mod -o tmp llvm --vector-width 2  --ir benchmark --run --instance-size 10 --repeat 5 passes --nmodl-ast
[NMODL] [info] :: Processing ../test.mod
...
[NMODL] [warning] :: RenameVisitor :: Renaming variable v at 23.14 to epilogue_v

VOID nrn_cur_hh(INSTANCE_STRUCT *mech){
    INTEGER id
    INTEGER node_id, nai_id, ion_nai_id
    DOUBLE v
    for(id = 0; id<mech->node_count-1; id = id+2) {
        node_id = mech->node_index[id]
        nai_id = mech->ion_nai_index[id]
        ion_nai_id = mech->ion_nai_index[id]
        v = mech->voltage[node_id]
        mech->nai[id] = mech->ion_nai[nai_id]
        mech->il[id] = mech->gl[id]*(v-mech->el[id])
        mech->ion_nai[ion_nai_id] = mech->nai[id]
    }
    INTEGER epilogue_node_id, epilogue_nai_id, epilogue_ion_nai_id
    DOUBLE epilogue_v
    for(; id<mech->node_count; id = id+1) {
        epilogue_node_id = mech->node_index[id]
        epilogue_nai_id = mech->ion_nai_index[id]
        epilogue_ion_nai_id = mech->ion_nai_index[id]
        epilogue_v = mech->voltage[epilogue_node_id]
        mech->nai[id] = mech->ion_nai[epilogue_nai_id]
        mech->il[id] = mech->gl[id]*(epilogue_v-mech->el[id])
        mech->ion_nai[epilogue_ion_nai_id] = mech->nai[id]
    }
}

[NMODL] [warning] :: RenameVisitor :: Renaming variable node_id at 1.21-27 to epilogue_node_id
...
[NMODL] [warning] :: RenameVisitor :: Renaming variable v at 1.21 to epilogue_v
VOID nrn_state_hh(INSTANCE_STRUCT *mech){
    INTEGER id
    INTEGER node_id, nai_id, ion_nai_id
    DOUBLE v
    for(id = 0; id<mech->node_count-1; id = id+2) {
        node_id = mech->node_index[id]
        nai_id = mech->ion_nai_index[id]
        ion_nai_id = mech->ion_nai_index[id]
        v = mech->voltage[node_id]
        mech->nai[id] = mech->ion_nai[nai_id]
        mech->m[id] = mech->minf[id]+mech->mtau[id]+0.1
        printf("%lf %d %lf", mech->m[id], ion_nai_id, mech->nai[id])
        mech->ion_nai[ion_nai_id] = mech->nai[id]
    }
    INTEGER epilogue_node_id, epilogue_nai_id, epilogue_ion_nai_id
    DOUBLE epilogue_v
    for(; id<mech->node_count; id = id+1) {
        epilogue_node_id = mech->node_index[id]
        epilogue_nai_id = mech->ion_nai_index[id]
        epilogue_ion_nai_id = mech->ion_nai_index[id]
        epilogue_v = mech->voltage[epilogue_node_id]
        mech->nai[id] = mech->ion_nai[epilogue_nai_id]
        mech->m[id] = mech->minf[id]+mech->mtau[id]+0.1
        printf("%lf %d %lf", mech->m[id], epilogue_ion_nai_id, mech->nai[id])
        mech->ion_nai[epilogue_ion_nai_id] = mech->nai[id]
    }
}

At the moment, nrn_cur looks similar to nrn_state and it will change slightly with additional code statements. But overall, it will have same operations as nrn_state i.e. no new operations than what you have already implemented for nrn_state. So it should be straightforward!

Let me know if you have any questions.

[georgemitenkov]

@pramodk great! I will see if any changes need to be made to the current pipeline tonight!

But from your example it seems like IR was generated and nothing failed☺️ (And scatter instructions are ignored until #643 is merged)

[pramodk]

Yeah, IR is generating fine.
Some of the tests were failing because we now have two codegen functions.I did minor changes and most of the things are looking OK.

[bbpbuildbot]

Logfiles from GitLab pipeline #47891 (:no_entry:) have been uploaded here!

Status and direct links:

• ✅ spack_setup
• ✅ build:intel
• ✅ build:gcc
• ⛔ test:intel
• ⛔ test:gcc

[pramodk]

I have rebased this PR now on latest llvm branch. By using hh.mod or simple test below:

NEURON {
    SUFFIX hh
    USEION na READ ena WRITE ina
    NONSPECIFIC_CURRENT il
    RANGE gnabar, gkbar, gl, el, gna, gk
    RANGE minf, hinf, ninf, mtau, htau, ntau
}

STATE {
    m h
}

ASSIGNED {
    v (mV)
    ena (mV)
    gna (S/cm2)
    ina (mA/cm2)
    minf hinf ninf
    mtau (ms) htau (ms) ntau (ms)
    gnabar
}

BREAKPOINT {
    SOLVE states METHOD cnexp
    gna = gnabar*m*m*m*h
    ina = gna*(v - ena)
}

DERIVATIVE states {
    m' =  (minf-m)/mtau
}

we get now below error:

./bin/nmodl test.mod llvm --ir
[NMODL] [info] :: Processing test.mod
[NMODL] [info] :: Running semantic analysis visitor
[NMODL] [info] :: Running symtab visitor
[NMODL] [info] :: Running CVode to cnexp 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 local variable rename visitor
[NMODL] [info] :: Running cnexp visitor
[NMODL] [info] :: Running C backend code generator
[NMODL] [info] :: Running LLVM backend code generator
[NMODL] [info] :: Running CodegenLLVMHelperVisitor
VOID nrn_cur_hh(INSTANCE_STRUCT *mech){
    INTEGER id
    INTEGER node_id, ena_id, ion_ina_id
    DOUBLE v
    for(id = 0; id<mech->node_count; id = id+1) {
        node_id = mech->node_index[id]
        ena_id = mech->ion_ena_index[id]
        ion_ina_id = mech->ion_ina_index[id]
        v = mech->voltage[node_id]
        mech->ena[id] = mech->ion_ena[ena_id]
        mech->gna[id] = mech->gnabar[id]*mech->m[id]*mech->m[id]*mech->m[id]*mech->h[id]
        mech->ina[id] = mech->gna[id]*(v-mech->ena[id])
        mech->ion_ina[ion_ina_id] += mech->ina[id]
    }
}
VOID nrn_state_hh(INSTANCE_STRUCT *mech){
    INTEGER id
    INTEGER node_id, ena_id
    DOUBLE v
    for(id = 0; id<mech->node_count; id = id+1) {
        node_id = mech->node_index[id]
        ena_id = mech->ion_ena_index[id]
        v = mech->voltage[node_id]
        mech->ena[id] = mech->ion_ena[ena_id]
        mech->m[id] = mech->m[id]+(1.0-exp(mech->dt*((((-1.0)))/mech->mtau[id])))*(-(((mech->minf[id]))/mech->mtau[id])/((((-1.0)))/mech->mtau[id])-mech->m[id])
    }
}

libc++abi: terminating with uncaught exception of type std::runtime_error:
              Error: only assignment is supported for CodegenAtomicStatement

The reason is that mech->ion_ina[ion_ina_id] += mech->ina[id] uses + with CodegenAtomicStatement.

@georgemitenkov: in this case, ion_ina_id is distinct (i.e. no conflict) for each iteration represented by id.

• So if I am not mistaken, we really don't need to create CodegenAtomicStatement here? I think rest of the code generation will be still correct, right? 🤔 cc: @iomaganaris
• Just to think bit ahead: assume that ion_ina_id has conflict i.e. within SIMD width, multiple ion_ina_id might have same value. In such case, any alternatives in LLVM IR to represent this? (In ISPC world there are atomics within vector width and I vaguely remember something in LLVM side)

There is #834 where we want to support atomic but that is for GPU backend only.

[bbpbuildbot]

Logfiles from GitLab pipeline #47986 (:no_entry:) have been uploaded here!

Status and direct links:

• ✅ spack_setup
• ✅ build:intel
• ✅ build:gcc
• ⛔ test:intel
• ⛔ test:gcc

[pramodk]

By the way, I am reading the code here:

nmodl/src/codegen/llvm/codegen_llvm_visitor.cpp

Line 486 in 6c3fe22

"Error: only assignment is supported for CodegenAtomicStatement\n");

but, addition of the form:

a[i] +=  b

can be written as

a[i] = a[i] + b

So the error we are throwing is not correct / consistent.

[bbpbuildbot]

Logfiles from GitLab pipeline #47996 (:no_entry:) have been uploaded here!

Status and direct links:

• ✅ spack_setup
• ✅ build:intel
• ✅ build:gcc
• ⛔ test:intel
• ⛔ test:gcc

[iomaganaris]

So if I am not mistaken, we really don't need to create CodegenAtomicStatement here? I think rest of the code generation will be still correct, right?

I think you're right about that.

The variables that should be updated atomically in the end is vec_rhs and vec_d I believe which are using node_id as index and currently are not implemented in the llvm backend

[bbpbuildbot]

Logfiles from GitLab pipeline #50520 (:white_check_mark:) have been uploaded here!

Status and direct links:

• ✅ spack_setup
• ✅ build:intel
• ✅ build:gcc
• ✅ test:intel
• ✅ test:gcc

[bbpbuildbot]

Logfiles from GitLab pipeline #50524 (:white_check_mark:) have been uploaded here!

Status and direct links:

• ✅ spack_setup
• ✅ build:intel
• ✅ build:gcc
• ✅ test:intel
• ✅ test:gcc

[pramodk]

@iomaganaris / @georgemitenkov : with latest fixes, all tests should pass. Now there are two kernels generated:

110: ; Function Attrs: nofree nounwind
110: define void @nrn_cur_hh(%__instance_var__type* noalias nocapture readonly %mech1) #0 {
110:   %mech = alloca %__instance_var__type*, align 8
110:   %id = alloca i32, align 4
110:   %node_id = alloca i32, align 4
110:   %ena_id = alloca i32, align 4
110:   %ek_id = alloca i32, align 4
110:   %ion_ina_id = alloca i32, align 4
110:   %ion_ik_id = alloca i32, align 4
110:   %v = alloca double, align 8
110:   store %__instance_var__type* %mech1, %__instance_var__type** %mech, align 8
110:   store i32 0, i32* %id, align 4
110:   br label %for.cond
110:
110: for.cond:                                         ; preds = %for.inc, %0
110:   %1 = load %__instance_var__type*, %__instance_var__type** %mech, align 8
110:   %2 = getelementptr inbounds %__instance_var__type, %__instance_var__type* %1, i32 0, i32 43
110:   %3 = load i32, i32* %2, align 4
110:   %4 = load i32, i32* %id, align 4
110:   %5 = icmp slt i32 %4, %3
110:   br i1 %5, label %for.body, label %for.exit, !llvm.loop !0
110:...
...
110:
110: ; Function Attrs: nofree nounwind
110: define void @nrn_state_hh(%__instance_var__type.0* noalias nocapture readonly %mech1) #0 {
110:   %mech = alloca %__instance_var__type.0*, align 8
110:   %id = alloca i32, align 4
110:   %node_id = alloca i32, align 4
110:   %ena_id = alloca i32, align 4
110:   %ek_id = alloca i32, align 4
110:   %v = alloca double, align 8
110:   store %__instance_var__type.0* %mech1, %__instance_var__type.0** %mech, align 8
110:   store i32 0, i32* %id, align 4
110:   br label %for.cond
110:
110: for.cond:                                         ; preds = %for.inc, %0
110:   %1 = load %__instance_var__type.0*, %__instance_var__type.0** %mech, align 8
110:   %2 = getelementptr inbounds %__instance_var__type.0, %__instance_var__type.0* %1, i32 0, i32 43
110:   %3 = load i32, i32* %2, align 4
110:   %4 = load i32, i32* %id, align 4
110:   %5 = icmp slt i32 %4, %3
110:   br i1 %5, label %for.body, label %for.exit, !llvm.loop !2

I have still one todo to take care of but that's not a blocker i.e. it just makes the code "computationally same" as neuron/coreneuron. I think you can already see if/how the nrn_cur can be called from benchmarking side. If desired, we can merge this PR and I can work in a separate branch/PR. Anyway, we will discuss it tomorrow.

[iomaganaris]

LGTM at the moment.
Benchmark is also running both nrn_cur and nrn_state kernels

[bbpbuildbot]

Logfiles from GitLab pipeline #51556 (:white_check_mark:) have been uploaded here!

Status and direct links:

• ✅ spack_setup
• ✅ build:intel
• ✅ build:gcc
• ✅ test:intel
• ✅ test:gcc

[pramodk]

@iomaganaris : I have completed nrn_cur implementation. Using vanilla mod file from coreneuron I see:

➜   ./bin/nmodl ~/workarena/repos/bbp/coreneuron/coreneuron/mechanism/mech/modfile/expsyn.mod llvm --ir passes

VOID nrn_cur_expsyn(INSTANCE_STRUCT *mech){
    INTEGER id
    INTEGER node_id, node_area_id
    DOUBLE v, g, rhs, v_org, current, mfactor
    for(id = 0; id<mech->node_count; id = id+1) {
        node_id = mech->node_index[id]
        node_area_id = mech->node_area_index[id]
        v = mech->voltage[node_id]
        v_org = v
        v = v+0.001
        {
            current = 0
            mech->i[id] = mech->g[id]*(v-mech->e[id])
            current = current+mech->i[id]
            mech->g[id] = current
        }
        v = v_org
        {
            current = 0
            mech->i[id] = mech->g[id]*(v-mech->e[id])
            current = current+mech->i[id]
            rhs = current
        }
        mech->g[id] = (mech->g[id]-rhs)/0.001
        mfactor = 1.e2/mech->node_area[node_area_id]
        mech->g[id] = mech->g[id]*mfactor
        rhs = rhs*mfactor
        mech->vec_rhs[node_id] = mech->vec_rhs[node_id]-rhs
        mech->vec_d[node_id] = mech->vec_d[node_id]+mech->g[id]
    }
}
VOID nrn_state_expsyn(INSTANCE_STRUCT *mech){
    INTEGER id
    INTEGER node_id
    DOUBLE v
    for(id = 0; id<mech->node_count; id = id+1) {
        node_id = mech->node_index[id]
        v = mech->voltage[node_id]
        mech->g[id] = mech->g[id]+(1.0-exp(mech->dt*((-1.0)/mech->tau[id])))*(-(0.0)/((-1.0)/mech->tau[id])-mech->g[id])
    }
}

➜  ./bin/nmodl ~/workarena/repos/bbp/coreneuron/coreneuron/mechanism/mech/modfile/hh.mod llvm --ir passes

VOID nrn_cur_hh(INSTANCE_STRUCT *mech){
    INTEGER id
    INTEGER node_id, ena_id, ek_id, ion_dinadv_id, ion_dikdv_id, ion_ina_id, ion_ik_id
    DOUBLE v, g, rhs, v_org, current, dina, dik
    for(id = 0; id<mech->node_count; id = id+1) {
        node_id = mech->node_index[id]
        ena_id = mech->ion_ena_index[id]
        ek_id = mech->ion_ek_index[id]
        ion_dinadv_id = mech->ion_dinadv_index[id]
        ion_dikdv_id = mech->ion_dikdv_index[id]
        ion_ina_id = mech->ion_ina_index[id]
        ion_ik_id = mech->ion_ik_index[id]
        v = mech->voltage[node_id]
        mech->ena[id] = mech->ion_ena[ena_id]
        mech->ek[id] = mech->ion_ek[ek_id]
        v_org = v
        v = v+0.001
        {
            current = 0
            mech->gna[id] = mech->gnabar[id]*mech->m[id]*mech->m[id]*mech->m[id]*mech->h[id]
            mech->ina[id] = mech->gna[id]*(v-mech->ena[id])
            mech->gk[id] = mech->gkbar[id]*mech->n[id]*mech->n[id]*mech->n[id]*mech->n[id]
            mech->ik[id] = mech->gk[id]*(v-mech->ek[id])
            mech->il[id] = mech->gl[id]*(v-mech->el[id])
            current = current+mech->il[id]
            current = current+mech->ina[id]
            current = current+mech->ik[id]
            g = current
        }
        dina = mech->ina[id]
        dik = mech->ik[id]
        v = v_org
        {
            current = 0
            mech->gna[id] = mech->gnabar[id]*mech->m[id]*mech->m[id]*mech->m[id]*mech->h[id]
            mech->ina[id] = mech->gna[id]*(v-mech->ena[id])
            mech->gk[id] = mech->gkbar[id]*mech->n[id]*mech->n[id]*mech->n[id]*mech->n[id]
            mech->ik[id] = mech->gk[id]*(v-mech->ek[id])
            mech->il[id] = mech->gl[id]*(v-mech->el[id])
            current = current+mech->il[id]
            current = current+mech->ina[id]
            current = current+mech->ik[id]
            rhs = current
        }
        g = (g-rhs)/0.001
        mech->ion_dinadv[ion_dinadv_id] = mech->ion_dinadv[ion_dinadv_id]+(dina-mech->ina[id])/0.001
        mech->ion_dikdv[ion_dikdv_id] = mech->ion_dikdv[ion_dikdv_id]+(dik-mech->ik[id])/0.001
        mech->ion_ina[ion_ina_id] = mech->ion_ina[ion_ina_id]+mech->ina[id]
        mech->ion_ik[ion_ik_id] = mech->ion_ik[ion_ik_id]+mech->ik[id]
        mech->vec_rhs[node_id] = mech->vec_rhs[node_id]-rhs
        mech->vec_d[node_id] = mech->vec_d[node_id]+g
    }
}
VOID nrn_state_hh(INSTANCE_STRUCT *mech){
    INTEGER id
    INTEGER node_id, ena_id, ek_id
    DOUBLE v
    for(id = 0; id<mech->node_count; id = id+1) {
        node_id = mech->node_index[id]
        ena_id = mech->ion_ena_index[id]
        ek_id = mech->ion_ek_index[id]
        v = mech->voltage[node_id]
        mech->ena[id] = mech->ion_ena[ena_id]
        mech->ek[id] = mech->ion_ek[ek_id]
        {
            DOUBLE alpha, beta, sum, q10, vtrap_in_0, vtrap_in_1, v_in_1
            v_in_1 = v
            UNITSOFF
            q10 = 3^((mech->celsius-6.3)/10)
            {
                DOUBLE x_in_0, y_in_0
                x_in_0 = -(v_in_1+40)
                y_in_0 = 10
                IF (fabs(x_in_0/y_in_0)<1e-6) {
                    vtrap_in_0 = y_in_0*(1-x_in_0/y_in_0/2)
                } ELSE {
                    vtrap_in_0 = x_in_0/(exp(x_in_0/y_in_0)-1)
                }
            }
            alpha = .1*vtrap_in_0
            beta = 4*exp(-(v_in_1+65)/18)
            sum = alpha+beta
            mech->mtau[id] = 1/(q10*sum)
            mech->minf[id] = alpha/sum
            alpha = .07*exp(-(v_in_1+65)/20)
            beta = 1/(exp(-(v_in_1+35)/10)+1)
            sum = alpha+beta
            mech->htau[id] = 1/(q10*sum)
            mech->hinf[id] = alpha/sum
            {
                DOUBLE x_in_1, y_in_1
                x_in_1 = -(v_in_1+55)
                y_in_1 = 10
                IF (fabs(x_in_1/y_in_1)<1e-6) {
                    vtrap_in_1 = y_in_1*(1-x_in_1/y_in_1/2)
                } ELSE {
                    vtrap_in_1 = x_in_1/(exp(x_in_1/y_in_1)-1)
                }
            }
            alpha = .01*vtrap_in_1
            beta = .125*exp(-(v_in_1+65)/80)
            sum = alpha+beta
            mech->ntau[id] = 1/(q10*sum)
            mech->ninf[id] = alpha/sum
        }
        mech->m[id] = mech->m[id]+(1.0-exp(mech->dt*((((-1.0)))/mech->mtau[id])))*(-(((mech->minf[id]))/mech->mtau[id])/((((-1.0)))/mech->mtau[id])-mech->m[id])
        mech->h[id] = mech->h[id]+(1.0-exp(mech->dt*((((-1.0)))/mech->htau[id])))*(-(((mech->hinf[id]))/mech->htau[id])/((((-1.0)))/mech->htau[id])-mech->h[id])
        mech->n[id] = mech->n[id]+(1.0-exp(mech->dt*((((-1.0)))/mech->ntau[id])))*(-(((mech->ninf[id]))/mech->ntau[id])/((((-1.0)))/mech->ntau[id])-mech->n[id])
    }
}

The only thing I wish to add is one test for synapse's BREAKPOINT block. So this is now ready for review.

[bbpbuildbot]

Logfiles from GitLab pipeline #51560 (:white_check_mark:) have been uploaded here!

Status and direct links:

• ✅ spack_setup
• ✅ build:intel
• ✅ build:gcc
• ✅ test:intel
• ✅ test:gcc

[georgemitenkov]

@pramodk I rebased atomic support on top of your PR and noticed that 2 tests fail in codegen_llvm_execution.cpp. This is because 2 kernels ("Simple scalar kernel", "Simple vectorised kernel") do not have i as assigned variable - and therefore helper visitor does not rename i to mech->i[id]. Could you add this change to that file as well?

[pramodk]

Could you add this change to that file as well?

Sure, will do!

[iomaganaris]

LGTM
As discussed offline we should just add an issue regarding the atomic reduction of the rhs and d variables
(I will approve once you add the tests)

[pramodk]

@iomaganaris : b9410dc add necessary test and addresses review comments.

[bbpbuildbot]

Logfiles from GitLab pipeline #51875 (:white_check_mark:) have been uploaded here!

Status and direct links:

• ✅ spack_setup
• ✅ build:intel
• ✅ build:gcc
• ✅ test:intel
• ✅ test:gcc

[iomaganaris]

LGTM after George's fix is merged

[bbpbuildbot]

Logfiles from GitLab pipeline #52089 (:white_check_mark:) have been uploaded here!

Status and direct links:

• ✅ spack_setup
• ✅ build:intel
• ✅ build:gcc
• ✅ test:intel
• ✅ test:gcc