llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
torch-mlir/test/Dialect/ATen/aten_batchnorm.mlir

25 lines
1.3 KiB
MLIR
Raw Normal View History

Add ATen Dialect (#16) This patch adds a dialect intended to be used as a frontend dialect to facilitate lowering from "A Tensor Library" in torch/pytorch. This patch includes several passes that are useful in conjuction with the dialect: --aten-layer-name: Generates layer names for each operation, which are not present in the original pytorch. --aten-to-std: Lower the ATen dialect into standard dialect function calls. --return-elimination-pass: convert functions (primarily the toplevel function) to pass return values by reference. This simplifies pytorch integration. --aten-op-report: generate a textual report about the model --liveness-report Future patches will implement actual integration with the pytorch jit to intercept and generates MLIR in this dialect, then lower the resulting MLIR into function calls through aten-layer-name -> aten-to-std -> return-elimination -> std-to-llvm. The result would then jitted using the LLVM jit, linked against a runtime library which makes calls back into pytorch to implement all the layers. Co-authored-by: Jeff Fifield <jeff.fifield@xilinx.com> Co-authored-by: Jeff Fifield <jeff.fifield@xilinx.com>
2020-08-13 10:28:04 +08:00
// RUN: npcomp-opt %s -aten-layer-name -aten-op-report |& FileCheck %s
// CHECK-LABEL: "L0-batch_norm-0": {
// CHECK-NEXT: "activation_in": 103320,
// CHECK-NEXT: "activation_out": 103320,
// CHECK-NEXT: "ops:*": 310206,
// CHECK-NEXT: "ops:+": 413280,
// CHECK-NEXT: "ops:-": 123,
// CHECK-NEXT: "ops:/": 123,
// CHECK-NEXT: "ops:sqrt": 123,
// CHECK-NEXT: "parameters_in": 246,
// CHECK-NEXT: "reads": 103566,
// CHECK-NEXT: "writes": 103320
module {
func @graph(%arg0: tensor<42x123x4x5xf32>, %arg1: tensor<123xf32>, %arg2: tensor<123xf32>, %arg3: tensor<123xf32>, %arg4: tensor<123xf32>, %arg5: tensor<?xi64>) -> tensor<42x123x4x5xf32> {
%0 = "aten.constant"() {type = "bool", value = 0 : i1} : () -> i1
%1 = "aten.constant"() {type = "f32", value = 1.000000e-01 : f32} : () -> f32
%2 = "aten.constant"() {type = "f32", value = 9.99999974E-6 : f32} : () -> f32
%3 = "aten.constant"() {type = "bool", value = 1 : i1} : () -> i1
%4:3 = "aten.batch_norm"(%arg0, %arg1, %arg2, %arg3, %arg4, %0, %1, %2, %3) : (tensor<42x123x4x5xf32>, tensor<123xf32>, tensor
<123xf32>, tensor<123xf32>, tensor<123xf32>, i1, f32, f32, i1) -> (tensor<42x123x4x5xf32>, tensor<123xf32>, tensor<123xf32>)
return %4#0 : tensor<42x123x4x5xf32>
}
}