ELECTR transpile proof of principle

[darcymason]

As per discussion on pymedphys/pymedphys#399 and then on nrc-cnrc/EGSnrc#658, looking at "transpiling" a test EGSnrc code to Python. A transpiler may be able to do it from Fortran, but here trying this test code just from Mortran.

To this point I've written a script, convert which does re.sub to transfer basic statements, comments, etc. to Python. Doing search/replace in an editor may have been possible, but with inevitable mistakes that would be discovered later, I thought it good to be able to 'replay' all steps - and can do that with the script starting from scratch every time.

The convert script converts subroutine ELECTR from egsnrc.mortran (stored in electr.mortran) to electr.py. For a while it also converted egsnrc.macros - at least some basic constants, and comments etc. to common.py

The end goal here is not optimal code, but something that can work and be hopefully be demonstrated equivalent to the mortran version through testing with the same random sequence.

The next steps:

• Convert loops (in progress)
• Convert remaining macros statements in electr.mortran into function calls
• Copy all constants and simple variables derived from egsnrc.macros to electr.py to act as globals
• As above, but for COMIN blocks.
• Comment out remaining unneeded (for this test) code
• set import statements for math functions as needed (sin) etc.
• Deal with gotos
• Correct indentation
• figure out 'interval'-based variables -> Python slices?
• Make mock functions to supply input, and for any processing needed outside this function
• continue to try to get to runnable Python code on single CPU
• [Ultimately, maybe not here, and probably not by me] try to get into jax form as per @SimonBiggs original suggestion

And some things to check:

• log and trig functions are in same units

[SimonBiggs]

I've got a friend coming over now for a couple of days, so likely won't get back to this until he goes home. Looking forward to diving in then though 🙂.

[SimonBiggs]

By the way... it's amazing just how much easier even syntax highlighting makes to reading the code.

[ftessier]

That is pretty awesome @darcymason!

[darcymason]

That is pretty awesome @darcymason!

Thanks, but the hard part is coming; We're on the wrong side of the Pareto Principle now 😉

And soon we'll have to start talking about how to mock in things like ausgab, howfar, etc. for the testing, and how to check that the code is even functioning correctly once it is working Python code. It may have some impact on the transpiling code.

[ftessier]

No doubt 😄 . For starters, I would leave ausgab and howfar as noops, since these are essentially callbacks to "user" space for the purpose of geometry definition and scoring. Or else one can write a simple geometry howfar, hownear and asugab (e.g., that for energy deposition in an infinite plate, as in the tutor1 mortran sample).

[darcymason]

So latest change needs some discussion. I've been thinking about the common blocks, and I don't see how parallelization can work with global variables. Each GPU needs its own its own variables for anything it is changing (e.g. NP for current position of particle on the stack).

As a suggested start, I've added a Particle class, which contains the common/stack variables (the ones that are arrays over NP). Instead of "passing in" (literally or by globals) a whole bunch of variables, at least these are in one nice package that is independent of other particles for parallelization.

I've also tried to assign the variable type conversion (e.g. Mortran $REAL-> Python float) all in one place, so they can more easily be changed as needed, e.g. to numpy or jax.numpy types.

The Particle class is just a start, there are probably a number of other COMMON variables that need a similar treatment. But many are perhaps "read-only" and can stay (physics data, medium info) - I'm not sure how smart jax is in recognizing those are okay - any thoughts, @SimonBiggs, when you have a chance?

[ftessier]

I don't want any globals in any EGSnrc rewrite, so I fully agree with your approach (see my comment about class encapsulation in the EGSnrc discussion). Years ago to implement a molecular dynamics code, I stuck with plain C (made sense in the 90s), but carried all the "globals" in a simulation struct (including function pointers, similar to objet methods!) and just passed the pointer around; worked like a charm, and was very efficient. At the time I could not achieve the same performance with C++. Trying to stick to pure functions (no modification of objects that live outside the function) is a useful mantra towards a good modular design.

[darcymason]

One other question: is there any easy way to find out which variables are input and which are output?

For example, in ELECTR, this doc is really nice, but I haven't seen much like it in other places:

    call msdist_pII
    (
      "Inputs
      eke,de,tustep,rhof,medium,qel,spin_effects,
      u(np),v(np),w(np),x(np),y(np),z(np),
      "Outputs
      uscat,vscat,wscat,xtrans,ytrans,ztrans,ustep
    );

Of course in Python, function arguments that are immutable (numbers, strings, etc) are not returned back to the caller - they need to be sent back in the return value of the function. So the distinction is quite important for this 'transpiler'.

[darcymason]

see my comment about class encapsulation in the EGSnrc discussion)

Ah, sorry, I thought I had scanned through that discussion, but I missed a lot, including your Particle class and @SimonBiggs comments on the parallelization of classes. I'll read through in more detail. I did see the Particles (plural) class, but I have trouble understanding how one can parallelize what happens in ELECTR to a parallel set of particles - given that there are a whole series of steps with branching conditions, and changes to different values, and quite often creation of new particles to track. Hence the need for some discussion, to educate me on these issues enough to make a reasonable first test.

[SimonBiggs]

but I have trouble understanding how one can parallelize what happens in ELECTR to a parallel set of particles - given that there are a whole series of steps with branching conditions, and changes to different values, and quite often creation of new particles to track

Yup. Here are some thoughts from a 2011 paper (GPUMCD: a new GPU-oriented monte carlo dose calculation platform, section II.C.3) for retaining vectorisation speed ups while being subject to divergence:

A notable difference between this work and the work of Jia et al. is the way the secondary particles are treated. In the work of Jia et al., a thread is responsible for its primary particle as well as every secondary particle it creates. This can be a major source of divergence because of the varying number of secondary particles created. On the other hand, GPUMCD does not immediately simulate secondary particles, but instead places them in their respective particle arrays. After a given pass of the simulation is over, the arrays are checked for newly created secondary particles and if secondary particles are found, they are simulated. This eliminates the divergence due to the different number of secondary particles per primary particles. Additionally, since Jia et al. used a single thread per primary particles, a thread may be responsible for the simulation of both electrons and photons, which can, in turn, be a major source of divergence. For example, at time t, thread A may be simulating a secondary electron, while thread B a secondary photon. In other words, they simulate both photons and electrons at the same time. Since both particles most likely have completely different code path, heavy serialization will occur. On the other hand, since GPUMCD stores secondary particles in arrays to be simulated later, no such divergence due to the electron-photon coupling occurs.

[darcymason]

On the other hand, GPUMCD does not immediately simulate secondary particles, but instead places them in their respective particle arrays

So I was reading more about jax from the links you gave, and was actually starting to think about the need (ultimately) to do just what they said there - and about separating particle types, to reduce the amount of branching.

All this adds up to a very different kind of code that we get from the transpile - for example, ELECTR will loop and simulate immediately any new electrons it generates (to save time in single CPU mode) - so that logic should go.

We'll have to have a good chat about all this... meanwhile I'll continue transpile work - in particular we'll still have to get physics data (or some kind of mock) and there are still some macro functions that haven't been converted yet. Some can probably be mocked, but some may need some close examination too to see if they are "pure" for vectorization. And this emphasizes again the need to know what variables are 'outputs' for each function.

[SimonBiggs]

You've probably gathered these answers by now, but just for completeness 🙂.

Each GPU needs its own its own variables for anything it is changing (e.g. NP for current position of particle on the stack).

JAX should be able to handle this memory management between the GPUs under the hood as long as we always provide it pure functions (see discussion at nrc-cnrc/EGSnrc#658 (reply in thread)). Also, static_argnums might help here as copied in at the end of this post.

I'm not sure how smart jax is in recognizing those are okay - any thoughts, @SimonBiggs, when you have a chance?

Yeah unfortunately no go there. From the jax.jit decorator docstring:

    fun: Function to be jitted. Should be a pure function, as side-effects may
      only be executed once. Its arguments and return value should be arrays,
      scalars, or (nested) standard Python containers (tuple/list/dict) thereof.
      Positional arguments indicated by ``static_argnums`` can be anything at
      all, provided they are hashable and have an equality operation defined.
      Static arguments are included as part of a compilation cache key, which is
      why hash and equality operators must be defined.

    static_argnums: An int or collection of ints specifying which positional
      arguments to treat as static (compile-time constant). Operations that only
      depend on static arguments will be constant-folded in Python (during
      tracing), and so the corresponding argument values can be any Python
      object. Static arguments should be hashable, meaning both ``__hash__`` and
      ``__eq__`` are implemented, and immutable. Calling the jitted function
      with different values for these constants will trigger recompilation. If
      the jitted function is called with fewer positional arguments than
      indicated by ``static_argnums`` then an error is raised. Arguments that
      are not arrays or containers thereof must be marked as static.
      Defaults to ().

That doesn't stop us using objects for syntactic sugar and correctness checks. But under the hood I believe it would be best that the data store backing of the objects be jax.DeviceArrays and any method on the class would offload the computationally expensive parts to an external pure function which just accepts jax arrays and standard Python objects.

[SimonBiggs]

Personally, for now, my preference is to avoid grouping things into custom classes altogether. Given we can't use those classes in the backend code, potentially it just adds extra overhead. And at the end of the day, those classes would be defined in a separate module altogether that is focused on API design.

Instead, to get the correctness benefits of custom classes, we can use mypy for typing checks, and define a "Particles" type that is just a restricted Python Dictionary. Given JAX can handle Python dictionaries this should work in our backend code, and that keeps it tested for correctness (at 'source code time') and keeps the syntactic sugar benefits.

I fleshed out that last idea over in PR #2. In particular:

https://github.com/SimonBiggs/egsnrc2py/blob/add-mypy-based-particles/prototyping/mypy_based_particles.py

[ftessier]

Indeed, any vectorized version of the code will rely on some version of the "Simultaneous Transport Of Particle Sets" (STOPS) idea introduced in vmc++ (see Kawrakow 2000 and Kawrakow and Fippel, 2000):

Groups of particles that have the same energy but different positions, directions, and possibly statistical weights, are simultaneously transported. This technique, referred to as STOPS - Simultaneous Transport Of Particle Sets, allows to re-use a variety of quantities such as mean-free-paths between interactions, sampled energy and direction changes, etc., and thus saves CPU time. The interaction type for each of the particles in the set is sampled separately based on the interaction probabilities in the medium the particle is in. Due to the very similar interaction probabilities for materials relevant for radiation therapy, in most cases all particles have the same interaction type. If, on the other side, different particles have different interactions selected, the set is split into sub-sets containing only interactions of the same type and each of the sub-sets transported recursively.

Then, as you mentioned, Sami's paper on GPUMCD provides great inspiration on the design of a vectorized version of the MC algorithm in general, and for grouping particles in particular. A word of caution, both vmc++ and GPUMCD were designed for particle transport in voxelized phantom for radiotherapy application, so not all of the trade-offs make sense when it comes to a reference general-purpose code such as EGSnrc.

Tong Xu from Carleton University has also implemented a GPU version of EGSnrc (but no publication yet as far as I am aware), where particles are saved in separate lists for each type of interaction.

@darcymason this implies that, as you said, the data structures and logic in the code will be markedly different, so transpiled code is unlikely to survive for something like ELECTR. On the other hand, your transpiler is useful to decipher other parts of the code, and I for one would still like to have a python version of EGSnrc on hand as a baseline reference implementation. And it is just fun, isn't it? 😉

[ftessier]

Personally, for now, my preference is to avoid grouping things into custom classes altogether. Given we can't use those classes in the backend code, potentially it just adds extra overhead. And at the end of the day, those classes would be defined in a separate module altogether that is focused on API design.

Agreed. If there is no way around flat arrays in the back end for vectorization (and I can't see any at this point), then let's adopt that mindset and then see how they can be packaged in API-ish form for the end user later on.

[darcymason]

On the other hand, your transpiler is useful to decipher other parts of the code, and I for one would still like to have a python version of EGSnrc on hand as a baseline reference implementation. And it is just fun, isn't it? 😉

Yes, it is fun ... and I've long thought it would be great to have a Python version of EGSnrc. It's easy to dismiss Python as being too slow, but as we know, computing changes over time. GPU vectorization didn't exist not so many years ago. Maybe in 5 or 10 years there will be something quite different. I certainly expect that with Python typing, compiled Python will get quite good, perhaps approaching Fortran or C speeds. With a little judicious use of numpy or vectorization of some expensive operations, a Python EGSnrc could closely approach past speeds.

So, is there a place for that? I could continue along the lines of what I've done, but expand to the "support code" as you mentioned. Perhaps it makes sense to separate the Python transpiling from the vectorization. @SimonBiggs, what do you think? I don't think I have the expertise (or the time to gain it) for the vectorization part. But perhaps if you did the proof-of-principle on that, I could continue with this transpiler, to have the start point "Pythonized" for different pieces, while keeping in mind the need for "pure" functions to help the vectorization. A kind of pipeline of Mortran/Fortran -> Python -> vectorized Python.

[ftessier]

There is a place for a python implementation, no doubt, if you enjoy working on that. I mean right now I don't see it becoming the production version of EGSnrc for a good while, until as you say JAX or other vectorization strategies pan out, or future processor and language development is strongly pythonic (which is not unlikely at this rate!). But having a python version of the current EGSnrc code is "a good thing": it provides a baseline reference, it allows for quickly implementing and testing a number of ideas, and can provide a really nice educational platform. Right now, it is very difficult for students (well, anyone really) to go in and modify the code without breaking it, owing to the globals and other obscure dependencies.

Moreover, with projects such as Cython, we could eventually compile the bottlenecks from python into C, while preserving a high-level pythonic view of the code (or we could link, for example, the existing egs++ library for geometries and sources into the python front-end). I often thought that a python code with C backend could provide the best of both worlds for EGSnrc: a natural and flexible programming language, and efficiency in the back end.

So at any rate, a transpiler from the current EGSnrc source code to python is valuable, in my view.

[SimonBiggs]

I certainly expect that with Python typing, compiled Python will get quite good, perhaps approaching Fortran or C speeds.

There is a huge benefit to have a reference transpilitation of EGSnrc no matter its speed. I personally hit a brick wall each time I try to read EGSnrc code. I could imagine the transpiler being included in the EGSnrc testing suite, so that each time EGSnrc code is changed, the transpiler can re-run, and then a range of tutor codes can be run on the EGSnrc fortran version, as well as the Python version and verified that they agree for every pull request.

A verifiably equivalent Python version of EGSnrc is massive for reasons such as readability and therefore a larger audience for the source code.

By trying to tackle vectorisation within the transpilation I think the complexity of the transpilation goes through the roof, and potentially it makes it difficult for it to ever become something that can be auto-tested. And, as you say, potentially down the line, the vectorisation approach we take may not end up being the fastest approach when computing changes. The "life-span" of the vectorised code may be on the order of 5 or so years until something else comes along that requires different opitimisations, the life-span of an autotested Python transpiler could last the lifetime of EGSnrc itself (as long as a condition of a mortran merge to the EGSnrc codebase is that the Python transpilations tests pass)...

But perhaps if you did the proof-of-principle on that, I could continue with this transpiler, to have the start point "Pythonized" for different pieces, while keeping in mind the need for "pure" functions to help the vectorization. A kind of pipeline of Mortran/Fortran -> Python -> vectorized Python.

I think that is a perfect division of labour.

[darcymason]

I could imagine the transpiler being included in the EGSnrc testing suite, so that each time EGSnrc code is changed, the transpiler can re-run

Unfortunately, I'm fairly certain transpiling will need manual intervention with the current state of the EGSnrc code. A big problem is that gotos can't be transpiled, they have to be converted to a Python branching structure of some kind which almost certainly needs human intervention to determine the logic. In some cases, a break statement might work, but automating that would be quite difficult. I imagine file I/O could be a problem, and the string handling looks very different, although maybe much of it can be automated.

However, having said all that, it probably wouldn't take much to manually update the transpiled code occasionally with changes to the base code, perhaps 'transpiling' the changed sections to get basic syntax differences, then hand-tune.

I've started playing with the transpyre library I mentioned before, on the Fortran produced by the Mortran run of tutor1, and it gives lots of error messages that are probably not helpful enough to work through without a lot of learning and digging. I succeeded with a very simple subroutine (a few lines) with a numeric variable passed in and assigned to another value in the function, and the transpyle command gave that back as a function returning None, when in ELECTR we would need that variable returned. So far, at least, my simple mostly-regex replacement seems less painful.

So at any rate, a transpiler from the current EGSnrc source code to python is valuable, in my view.

I will keep playing around with this, see what the scope might be and how far I can get - and try to build something testable.

[ftessier]

I see the transpiler more as a "one-off" effort, helping to bootstrap a python implementation, rather than something that would be used dynamically as EGSnrc evolves. EGSnrc must evolve OUT of Mortran and Fortran. In fact, there is no significant code evolution in the backend anymore (I mean in terms of number of code lines), mostly bug fixes and the occasional improvement. The EGSnrc structure prevents further development in my opinion. So I suggest not sweating over making the transpiler completely automatic. Manual intervention is fine: I think it would be sufficient to record manual tweaks as git commits which could be replayed on top of a new transpiled version, if need be.

[SimonBiggs]

I see the transpiler more as a "one-off" effort, helping to bootstrap a python implementation

Awesome, so essentially once done the Python transpiled version will become the reference EGSnrc version.

[ftessier]

Awesome, so essentially once done the Python transpiled version will become the reference EGSnrc version.

At this point I would say this is a long shot (efficiency, verification, validation, community uptake and confidence, etc.), but definitely worth trying! I support it, and will help. An efficient python implementation (on par with C or Fortran) would certainly gain lots of traction if it is well-vetted at every stage of development, and if the code remains clear and manageable despite vectorization. I am still considering C++ or Rust implementations though. Maybe we try them all!? 🤓

[SimonBiggs]

I am still considering C++ or Rust implementations though. Maybe we try them all!? 🤓

All the things 🙂

I have been liking this parallel version of the prototyping examples that you've been producing. It has been helpful to know what time is needed to be beaten. I'm not sure if you'd gain any speed improvements if you were to also be writing those examples in rust.

An efficient python implementation (on par with C or Fortran)

Would something that is within a factor of ~2-3 slower for CPU usage, but potentially at least an order of magnitude faster on GPU count as on-par?

I support it, and will help.

🙂

[ftessier]

😂 I think within an order of magnitude slower on CPU and at least one order of magnitude faster on a typical GPU is where I would put the on-par bar (in the sense that the effort is worth it). I am currently testing the C++, plain python and JAX using perf, as suggested by @crcrewso: it is very instructive in the limits of wall-clock timing! I will share those timing soon.

[darcymason]

An update:

The EGSnrc structure prevents further development in my opinion.

Interestingly, as I go to post an update here to what I've been working on, the above quote is a really good synopsis of what I've learned.

A big part of that is the global variables problem, and also the whole macro concept - very flexible in Mortran, but complicated to port to a different language. Its necessary to get through dealing with those before much can be done in any pathway. And there are very often macros within macros (replacement text includes macros in it, and even sometimes a 'language' within the macros of IF/ELSE blocks etc.).

Add on top of that, I'm realizing how much a language from that time is limited compared with modern ones. But for now, will try mostly for a direct mapping and not worry about optimizations (that's pyegsnrc's angle, hopefully we can meet in the middle).

On the macro side, I'm turning macros into:
- a "parameters" file for what look like constants or parameters.
- in-lining of shorter macro replacements (like Mortran does for speed), but sandwiched in comments so they can easily be found and modified if needed.
- separate functions for longer replacements, to try to keep the code more readable. These will need function arguments later
- many function calls are empty in EGSnrc, to give the user callbacks to adjust the physics or record information. Those are set to None, but still called if not None (a None check is much faster than a function call), to leave those options available.

Because it was interesting, I looked again at formal "transpilers" - the transpyle library, "lark", ANTLR, and played around with these, with Fortran as the source code
- for transpyle, With "real" EGSnrc code, I went down a number of paths without being able to get it to work, and there are many pre-processing steps still necessary. Hand-editing would always be required anyway.
- I explored creating a Mortran grammar, using Python libary "lark". Got a fair way, able to correctly code math expressions, variables, subroutines, IF/ELSEIF/ELSE blocks, and LOOP. Painstaking to get it right, though. And turning that back into Python is tricky, but mostly do-able, I think. However, there is much more that would need to be done.

At this stage, I paused the grammar approach for now, go back to some regex-based conversion and hand-editing to learn more about the scope of the conversion and other issues yet to be discovered. At some point hand-editing would actually be less time overall than coding automatic transpiling.

tl;dr: I'm continuing to automate much of the macro conversion, but that's almost done for this round - almost to the point of hand-edits for this pilot (continuing on subroutine ELECTR, as had already started there). It's not too much different than my checklist at the top, but I've gone a little deeper into converting functions and macros.

[ftessier]

Thank you for the detailed update and comments on the rigidity of Mortran, from your "inside" perspective! One has to admire---and I do---the cleverness and prescience of the Mortran authors in implementing macros back in the 1970s. Today these would (will) be function pointers, or just functions in languages where functions are first-class citizens, as in python. As you have mentioned, nested macros implies that it is not trivial to reproduce EGSnrc macros, even with modern languages. I have tried at times...

In fact, I don't really understand how Mortran is able to do what it does, i.e., keep replacing macros until there are none left (it seems it could get stuck in a loop, but it does not!). This goes back to the powerful in the 1991 comment in the EGS code I quoted before: "And yes, Mortran MUST DIE but we have to put something as powerful in its place. Any volunteers?". Not sure this was a reasonable ask in 1991.

It is important (at least in a first ditch effort) to preserve this "replacement" ability, even for empty (None) macros. The power of macros in a toolkit such as EGSnrc is that it allows users to redefine them on the application side, without modifying the pristine source code. Advanced applications such as BEAMnrc and DOSXYZnrc use that to extend EGSnrc for their specific purposes. We want to keep this flexibility, i.e., we want EGSnrc to remain a toolkit with research in mind, where for example I can turn off the Compton interaction if I want, etc.

Of course, one could opt to modify the Mortran interpreter (transpiler!) itself to output python code instead of Fortran... Ha ha ha. 😉

At any rate, thank you for your interest and effort into transpiling the code: this is probably the best first step for a rewrite, bound to provide insight into the inner working of EGSnrc. In fact, you already know more than I do, no doubt!

[darcymason]

In fact, I don't really understand how Mortran is able to do what it does, i.e., keep replacing macros until there are none left (it seems it could get stuck in a loop, but it does not!)

I think I have my head around that now - the trick is to apply a substitution, then run all substitutions on that value, repeatedly, until the value does not change. The code here (currently in a PR) is close to being able to do such a full replace - in fact it is trying but some parts are not working yet, and I haven't got all types of macros being applied yet in the first place. However, this would still exclude the replacements with 'language structures' as I mentioned before - I've started hand-coding those - there is a pretty easy hand-edit translation to a re.sub replace function rather than string, which can do the substitutions with Python f-string quite easily. However, those complex macros then become fixed - based on the current snapshot of egsnrc.macros in this repo. I'm assuming no one probably changes those ones any more. I don't think we can keep all the original flexibility, but I'm trying to keep as much as reasonable.

Of course, one could opt to modify the Mortran interpreter (transpiler!) itself to output python code instead of Fortran... Ha ha ha. 😉

Yeah, that would be a lot of fun (sarcasm) - I've looked at the code for Mortran itself (written in Fortran of course; this whole thing starts to really warp ones ability to meta-think). I have to give kudos to what they were able to do, but it is completely incomprehensible, it seems to be just a bunch of replacement and swapping of numbers in arrays - I'm guessing many of those numbers are actually the text of the commands (like 'REPLACE') or the substitutions?

At any rate, thank you for your interest and effort into transpiling the code: this is probably the best first step for a rewrite, bound to provide insight into the inner working of EGSnrc.

I've gone far enough to know that it is an achievable goal at least (achievable as in technically possible; to do the entire code-base will be a long haul, though). At the very least, code in Python would be much more amenable to help guide translation to something else later.

In fact, you already know more than I do, no doubt!

I very much doubt that 😄. I'll be coming back soon with a number of questions, problems, help in importing data to set up test functions, etc.

[darcymason]

Okay, after hand-edits, I have a working electr, mostly in Python.

See discussion at darcymason/egsnrc#7