Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save matthias-springer/5020adf5ccdf43dcdc6ffd828b93a2a7 to your computer and use it in GitHub Desktop.
Save matthias-springer/5020adf5ccdf43dcdc6ffd828b93a2a7 to your computer and use it in GitHub Desktop.
// RUN: mlir-opt %s -test-transform-dialect-interpreter="debug-payload-root-tag=payload"
module attributes {transform.target_tag="payload"} {
// Batched TOSA matrix multiplication. %A and %B are the inputs, %C is the
// output.
func.func @test_matmul(%A: memref<1x17x19xf32>, %B: memref<1x19x29xf32>,
%C: memref<1x17x29xf32>) {
%A_tensor = bufferization.to_tensor %A restrict : memref<1x17x19xf32>
%B_tensor = bufferization.to_tensor %B restrict : memref<1x19x29xf32>
%0 = tosa.matmul %A_tensor, %B_tensor
: (tensor<1x17x19xf32>, tensor<1x19x29xf32>) -> tensor<1x17x29xf32>
bufferization.materialize_in_destination %0 in restrict writable %C
: (tensor<1x17x29xf32>, memref<1x17x29xf32>) -> ()
return
}
}
transform.sequence failures(propagate) {
^bb0(%module1: !transform.any_op):
// Convert TOSA to Linalg
%func1 = transform.structured.match ops{["func.func"]} in %module1
: (!transform.any_op) -> !transform.any_op
%func2 = transform.apply_registered_pass "tosa-to-linalg-named" to %func1
: (!transform.any_op) -> !transform.any_op
// Canonicalize + CSE
transform.apply_patterns to %func2 {
transform.apply_patterns.canonicalization
} {apply_cse} : !transform.any_op
// Eliminate empty tensors
transform.bufferization.eliminate_empty_tensors %func2 : !transform.any_op
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment