20 lines
1.4 KiB
MLIR
20 lines
1.4 KiB
MLIR
// RUN: mlir-hlo-opt -disc-ral-inject-execution-context=entry-func-name=test \
|
|
// RUN: -canonicalize %s -o - | FileCheck %s
|
|
|
|
// CHECK-LABEL: func @test
|
|
// CHECK-SAME: (%[[CTX:.*]]: !disc_ral.context) {
|
|
func @test(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>,
|
|
%arg2: memref<?x?xf32>, %arg3: memref<?x?xf32>) -> (memref<?x?xf32>, memref<?x?xf32>) {
|
|
// %[[T0:.*]] = "disc_ral.recv_input"(%[[CTX]], %c0) : (!disc_ral.context, index) -> memref<?x?xf32>
|
|
// %[[T1:.*]] = "disc_ral.recv_input"(%[[CTX]], %c1) : (!disc_ral.context, index) -> memref<?x?xf32>
|
|
// %[[T2:.*]] = "disc_ral.recv_input"(%[[CTX]], %c2) : (!disc_ral.context, index) -> memref<?x?xf32>
|
|
// %[[T3:.*]] = "disc_ral.recv_input"(%[[CTX]], %c3) : (!disc_ral.context, index) -> memref<?x?xf32>
|
|
// "lmhlo.abs"(%[[T0]], %[[T1]]) : (memref<?x?xf32>, memref<?x?xf32>) -> ()
|
|
// "lmhlo.add"(%[[T1]], %[[T2]], %[[T3]]) : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>) -> ()
|
|
// "disc_ral.send_output"(%[[CTX]], %c0, %[[T0]]) : (!disc_ral.context, index, memref<?x?xf32>) -> ()
|
|
// "disc_ral.send_output"(%[[CTX]], %c1, %[[T3]]) : (!disc_ral.context, index, memref<?x?xf32>) -> ()
|
|
"lmhlo.abs"(%arg0, %arg1) : (memref<?x?xf32>, memref<?x?xf32>) -> ()
|
|
"lmhlo.add"(%arg1, %arg2, %arg3) : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>) -> ()
|
|
return %arg1, %arg3 : memref<?x?xf32>, memref<?x?xf32>
|
|
}
|