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

[ONNX][TORCH-MLIR] LayerNorm (#2716) 

Layer Normalization using the torch.aten.native_layer_norm 

https://github.com/nod-ai/SHARK-Turbine/issues/325
pull/2772/head snapshot-20240111.1080
Andreas Falkenberg 2024-01-11 00:57:04 -08:00 committed by GitHub
parent 0860c41ee2
commit 5862854bc8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 56 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp
View File

@ -410,6 +410,49 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
} }
return failure(); return failure();
}); });
patterns.onOp("LayerNormalization", 17,
[](OpBinder binder, ConversionPatternRewriter &rewriter) {
Torch::ValueTensorType Y_type;
Torch::ValueTensorType Mean_type;
Torch::ValueTensorType InvStdDev_type;
Value X;
Value Scale;
Value B;
int64_t axis;
float epsilon;
int64_t stash_type;
if (binder.tensorOperandAtIndex(X, 0) ||
binder.tensorOperandAtIndex(Scale, 1) ||
binder.tensorOperandAtIndex(B, 2) ||
binder.tensorResultTypeAtIndex(Y_type, 0) ||
binder.tensorResultTypeAtIndex(Mean_type, 1) ||
binder.tensorResultTypeAtIndex(InvStdDev_type, 2) ||
binder.s64IntegerAttr(axis, "axis", -1) ||
binder.f32FloatAttr(epsilon, "epsilon", 0.00001) ||
binder.s64IntegerAttr(stash_type, "stash_type", 1))
return failure();
Value constEpsilon = rewriter.create<Torch::ConstantFloatOp>(
binder.getLoc(), rewriter.getType<Torch::FloatType>(),
rewriter.getF64FloatAttr(epsilon));
unsigned rank = 1;
if(std::optional<unsigned> maybeRank = Torch::getTensorRank(X))
rank = *maybeRank;
SmallVector<Value> normalized;
axis = Torch::toPositiveDim(axis, rank);
auto X_type = X.getType().cast<Torch::ValueTensorType>();
ArrayRef<int64_t> X_shape = X_type.getSizes();
for (int64_t n = axis; n < rank ; n++) {
normalized.push_back(rewriter.create<Torch::ConstantIntOp>(
binder.getLoc(), rewriter.getI64IntegerAttr(X_shape[n])));
}
Value normalized_shape = rewriter.create<Torch::PrimListConstructOp>(
binder.getLoc(),
Torch::ListType::get(Torch::IntType::get(binder.op->getContext())),
normalized);
rewriter.replaceOpWithNewOp<Torch::AtenNativeLayerNormOp>(
binder.op, Y_type, Mean_type, InvStdDev_type, X, normalized_shape, Scale, B, constEpsilon);
return success();
});
patterns.onOp("LeakyRelu", 1, patterns.onOp("LeakyRelu", 1,
[](OpBinder binder, ConversionPatternRewriter &rewriter) { [](OpBinder binder, ConversionPatternRewriter &rewriter) {
Torch::ValueTensorType resultType; Torch::ValueTensorType resultType;

13
test/Conversion/TorchOnnxToTorch/simple_ops_g_to_p.mlir
View File

@ -116,6 +116,19 @@ func.func @test_gemm_alpha_beta(%arg0: !torch.vtensor<[3,5],f32>, %arg1: !torch.
// ----- // -----
// CHECK-LABEL : func.func @test_layer_norm
func.func @test_layer_norm(%arg0: !torch.vtensor<[3,4],f32>, %arg1: !torch.vtensor<[3,4],f32>, %arg2: !torch.vtensor<[3,4],f32>) -> (!torch.vtensor<[3,4], f32>, !torch.vtensor<[1,1],f32>, !torch.vtensor<[1,1],f32>)
attributes {torch.onnx_meta.ir_version = 6 : si64, torch.onnx_meta.opset_version = 17 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
// CHECK: %int3 = torch.constant.int 3
// CHECK: %int4 = torch.constant.int 4
// CHECK: %0 = torch.prim.ListConstruct %int3, %int4 : (!torch.int, !torch.int) -> !torch.list<int>
// CHECK: %result0, %result1, %result2 = torch.aten.native_layer_norm %arg0, %0, %arg1, %arg2
%0:3 = torch.operator "onnx.LayerNormalization"(%arg0, %arg1, %arg2) {torch.onnx.axis = 0 : si64} : (!torch.vtensor<[3,4],f32>, !torch.vtensor<[3,4],f32>, !torch.vtensor<[3,4],f32>) -> (!torch.vtensor<[3,4],f32>, !torch.vtensor<[1,1],f32>, !torch.vtensor<[1,1],f32>)
return %0#0, %0#1, %0#2 : !torch.vtensor<[3,4],f32>, !torch.vtensor<[1,1],f32>, !torch.vtensor<[1,1],f32>
}
// -----
// CHECK-LABEL: func.func @test_leaky_relu // CHECK-LABEL: func.func @test_leaky_relu
func.func @test_leaky_relu(%arg0: !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],f32> attributes {torch.onnx_meta.opset_version = 16 : si64} { func.func @test_leaky_relu(%arg0: !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],f32> attributes {torch.onnx_meta.opset_version = 16 : si64} {
// CHECK-DAG: %[[F2:.+]] = torch.constant.float 2 // CHECK-DAG: %[[F2:.+]] = torch.constant.float 2