Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save matthias-springer/b3f40d1667c977c29a76cc7a469cc1a0 to your computer and use it in GitHub Desktop.
Save matthias-springer/b3f40d1667c977c29a76cc7a469cc1a0 to your computer and use it in GitHub Desktop.
// RUN: mlir-opt %s -test-transform-dialect-interpreter="debug-payload-root-tag=payload" -transform-dialect-drop-schedule
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
// Eliminate empty tensors
transform.bufferization.eliminate_empty_tensors %func2 : !transform.any_op
// Bufferize
// %module2 = transform.bufferization.one_shot_bufferize %module1
// : (!transform.any_op) -> !transform.any_op
// Canonicalize + CSE
// transform.apply_patterns to %module2 {
// transform.apply_patterns.canonicalization
// } {apply_cse} : !transform.any_op
// Deallocate buffers. That's a no-op in this example because no temporary
// buffers are allocated.
// %module3 = transform.apply_registered_pass
// "buffer-deallocation-pipeline" to %module2
// : (!transform.any_op) -> !transform.any_op
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment