// Graph-local structured input for loom-lower-graph-memory.
// Each module holds one or two dataflow.graph definitions whose bodies are
// source-sequential scf.while loops over normalized scalar memref.load /
// memref.store leaves on graph memory inputs (canonical roots).

start: {new NLOOP = random.randint(1, 2); new L = 0}
       'module {\n' graph_list '}\n';

graph_list: (L < NLOOP) graph_def {L += 1} graph_list
          | (L == NLOOP) '';

graph_def: {new NARGS = random.choice([2, 3]); new LIMIT = random.randint(2, 6); new INNER_STORE = random.choice([0, 1]); new EXTRA_LOAD = random.choice([0, 1]); new USE_ARG_BOUND = random.choice([0, 1]); new AFTER_MODE = random.choice([0, 1, 2])}
           'dataflow.graph private @while_case' graph_index
           '(%start: none, %limit: i32, %input: memref<?xi32>, %output: memref<?xi32>) -> ()\n'
           '    attributes {input_segments = array<i32: 1, 0, 2>,\n'
           '                result_segments = array<i32: 0, 0, 0>} {\n'
           '  %zero = arith.constant 0 : i32\n'
           '  %one = arith.constant 1 : i32\n'
           '  %bound = arith.constant ' limit_text ' : i32\n'
           '  %res:' nargs_text ' = scf.while (' init_list ') : (' type_list ') -> (' type_list ') {\n'
           before_region
           '  } do {\n'
           '  ^bb0(' after_args '):\n'
           after_body
           '  }\n'
           exit_uses
           '  dataflow.graph.return %start : none\n'
           '}\n';

graph_index: [str(L)];

limit_text: [str(LIMIT)];

nargs_text: [str(NARGS)];

init_list: (NARGS == 2) '%s0 = %zero, %s1 = %zero'
         | (NARGS == 3) '%s0 = %zero, %s1 = %zero, %s2 = %one';

type_list: (NARGS == 2) 'i32, i32'
         | (NARGS == 3) 'i32, i32, i32';

after_args: (NARGS == 2) '%b0: i32, %b1: i32'
          | (NARGS == 3) '%b0: i32, %b1: i32, %b2: i32';

// The after region either forwards its block values, recomputes one of them,
// or performs its own memory access before yielding.
after_body: after_statements '    scf.yield ' after_values ' : ' type_list '\n';

after_statements: (AFTER_MODE == 0) ''
                | (AFTER_MODE == 1) '    %d1 = arith.addi %b1, %b0 : i32\n'
                | (AFTER_MODE == 2) '    %bidx = arith.index_cast %b0 : i32 to index\n'
                                    '    memref.store %b1, %output[%bidx] : memref<?xi32>\n';

after_values: (AFTER_MODE == 1) '%b0, %d1' third_after_value
            | (AFTER_MODE == 0) '%b0, %b1' third_after_value
            | (AFTER_MODE == 2) '%b0, %b1' third_after_value;

third_after_value: (NARGS == 3) ', %b2'
                 | (NARGS == 2) '';

// The before region is the recursively lowered region of the while loop: it
// reads memory, computes the next state, and ends in scf.condition.
before_region: '    %idx0 = arith.index_cast %s0 : i32 to index\n'
               '    %v0 = memref.load %input[%idx0] : memref<?xi32>\n'
               '    %n0 = arith.addi %s0, %one : i32\n'
               second_load
               '    %n1 = arith.addi %s1, ' accumulated ' : i32\n'
               third_state
               inner_store
               '    %c = arith.cmpi slt, %n0, ' bound_value ' : i32\n'
               '    scf.condition(%c) ' condition_args ' : ' type_list '\n';

second_load: (EXTRA_LOAD == 1) '    %idx1 = arith.index_cast %n0 : i32 to index\n'
                               '    %v1 = memref.load %input[%idx1] : memref<?xi32>\n'
                               '    %vsum = arith.addi %v0, %v1 : i32\n'
           | (EXTRA_LOAD == 0) '';

accumulated: (EXTRA_LOAD == 1) '%vsum'
           | (EXTRA_LOAD == 0) '%v0';

third_state: (NARGS == 3) '    %n2 = arith.addi %s2, %n0 : i32\n'
           | (NARGS == 2) '';

inner_store: (INNER_STORE == 1) '    memref.store %n1, %output[%idx0] : memref<?xi32>\n'
           | (INNER_STORE == 0) '';

bound_value: (USE_ARG_BOUND == 1) '%limit'
           | (USE_ARG_BOUND == 0) '%bound';

condition_args: (NARGS == 2) '%n0, %n1'
              | (NARGS == 3) '%n0, %n1, %n2';

// Every while result is consumed after the loop, so the loop-exit lanes are
// observable in the lowered graph.
exit_uses: '  %oidx0 = arith.index_cast %zero : i32 to index\n'
           '  memref.store %res#0, %output[%oidx0] : memref<?xi32>\n'
           '  %oidx1 = arith.index_cast %one : i32 to index\n'
           '  memref.store %res#1, %output[%oidx1] : memref<?xi32>\n'
           third_exit_use;

third_exit_use: (NARGS == 3) '  %otwo = arith.constant 2 : i32\n'
                             '  %oidx2 = arith.index_cast %otwo : i32 to index\n'
                             '  memref.store %res#2, %output[%oidx2] : memref<?xi32>\n'
              | (NARGS == 2) '';