xdslproject/linalg-to-stream

linalg to Stream

Input: linalg dialect .mlir file

Output:

1. annotated linalg dialect (linalg generic operations that can be processed by stream are annotated with unique ids)
2. a python dictionary saved to the filename workload.py valid as stream input

set up

Install Dependencies

pip install -r requirements.txt

running the tool

python xdsl_opt_main.py tests/matmul.mlir -p linalg-to-stream

limitations

• Currently tool can only take in a single linalg generic operation
• The linalg generic operation must be a matrix multiply

future work

• handle multiple linalg generic operations, assigning a unique id to each, which is then added as as attribute to the mlir operation

• figure out dependencies between linalg generic operations, and record this relationship in the workload objects using the unique ids

• handle case where linalg generic has more than three operands (we are assuming the first two operands are inputs, and the last operand is an output) Quantized operations have more than two inputs, and we would like to support these.

  Example of a quantized operation we want to support:

  func.func @simple_matmul(%arg0: memref<16x16xi8>, %arg1: memref<16x16xi8, strided<[1, 16], offset:0>>, %arg2: memref<16x16xi32>) {
  %c0_i32 = arith.constant 0 : i32
  linalg.quantized_matmul ins(%arg0, %arg1, %c0_i32, %c0_i32 : memref<16x16xi8>, memref<16x16xi8, strided<[1, 16], offset:0>>, i32, i32) outs(%arg2 : memref<16x16xi32>)
  return
  }
  

  gobolt.org MLIR opt (trunk) --linalg-generalize-named-ops --mlir-print-local-scope --mlir-print-op-generic

  "builtin.module"() ({
    "func.func"() <{function_type = (memref<16x16xi8>, memref<16x16xi8, strided<[1, 16]>>, memref<16x16xi32>) -> (), sym_name = "simple_matmul"}> ({
    ^bb0(%arg0: memref<16x16xi8>, %arg1: memref<16x16xi8, strided<[1, 16]>>, %arg2: memref<16x16xi32>):
      %0 = "arith.constant"() <{value = 0 : i32}> : () -> i32
      "linalg.generic"(%arg0, %arg1, %0, %0, %arg2) <{indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> ()>, affine_map<(d0, d1, d2) -> ()>, affine_map<(d0, d1, d2) -> (d0, d1)>], iterator_types = [#linalg.iterator_type<parallel>, #linalg.iterator_type<parallel>, #linalg.iterator_type<reduction>], operandSegmentSizes = array<i32: 4, 1>}> ({
      ^bb0(%arg3: i8, %arg4: i8, %arg5: i32, %arg6: i32, %arg7: i32):
        %1 = "arith.extsi"(%arg3) : (i8) -> i32
        %2 = "arith.subi"(%1, %arg5) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
        %3 = "arith.extsi"(%arg4) : (i8) -> i32
        %4 = "arith.subi"(%3, %arg6) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
        %5 = "arith.muli"(%2, %4) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
        %6 = "arith.addi"(%arg7, %5) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
        "linalg.yield"(%6) : (i32) -> ()
      }) : (memref<16x16xi8>, memref<16x16xi8, strided<[1, 16]>>, i32, i32, memref<16x16xi32>) -> ()
      "func.return"() : () -> ()
    }) : () -> ()
  }) : () -> ()
  

Note: None of the MLIR matrix operations take padding as input, so we know all the workloads produced from linalg should have 0 padding.

Tests To Make:

• vector x vector ; elementwise multiplication
• conv2D
• Conv 1D
• Depthwise Conv2D*
• Pointwise Conv2D
• Matrix-vector multi.
• matrix-martix multiply

feed output of tool into zigzag

python run_zigzag.py 

feed output of tool into stream (need to fix)

python run_stream.py

Errors we're getting documented here

linalg-to-stream's Issues

Error in loop_dim_size conversion

for inputs.matmul2.mlir,

"builtin.module"() ({
  %0:3 = "test.op"() : () -> (memref<5x7xf32>, memref<7x6xf32>, memref<5x6xf32>)
  "linalg.generic"(%0#0, %0#1, %0#2) <{indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>], iterator_types = [#linalg.iterator_type<parallel>, #linalg.iterator_type<parallel>, #linalg.iterator_type<reduction>], operandSegmentSizes = array<i32: 2, 1>}> ({
  ^bb0(%arg0: f32, %arg1: f32, %arg2: f32):
    %1 = "arith.mulf"(%arg0, %arg1) <{fastmath = #arith.fastmath<none>}> : (f32, f32) -> f32
    %2 = "arith.addf"(%arg2, %1) <{fastmath = #arith.fastmath<none>}> : (f32, f32) -> f32
    "linalg.yield"(%2) : (f32) -> ()
  }) : (memref<5x7xf32>, memref<7x6xf32>, memref<5x6xf32>) -> ()
}) : () -> ()

we get the following workload,

workload = {
    0: {
        "operator_type": "default",
        "equation": "O[d0][d1] += I[d0][d2] * W[d2][d1]",
        "dimension_relations": [],
        "loop_dim_size": {"D0": 5, "D1": 7, "D2": 7},
        "operand_precision": {"O": 32, "O_final": 32, "W": 32, "I": 32},
        "operand_source": {"W": [], "I": []},
        "constant_operands": ["I", "W"],
        "padding": {},
    }
}

Which I believe is incorrect.
Shouldn't the loop_dim_size for D1 be 6 here?
Something is going wrong with the way loop_dim_size fields are parsed.

Add end to end testing CI

At the very least, a simple shell script with a diff.

Export workload as a json file

It would be nice to export the workload in a json file (we can format that, diff etc) and with a filename of our choice.
Creating a .py module should fall on the run_* scripts.

This will allow us to potentially store and create tests based on the expected json files, and also run multiple tests for different kernels and sequences of kernels without having to resort to moving workload.py tricks.