fix layer_norm half precision

examples/PyTorch/Inference/CUDA/BERT/main.py

@@ -263,17 +263,17 @@ def run():
     # Run naive torch.
     print("Naive PyTorch.")
     model = bert_large_amp
-    evaluate_torch(model, inputs)
+    # evaluate_torch(model, inputs)
 
     if utils.torch_version_number() >= utils.parse_version("1.14.0"):
         print("BladeDISC PyTorch 2.0 Optimization.")
-        evaluate_torch(torch.compile(bert_large_amp, backend="disc"), inputs)
+        evaluate_torch(torch.compile(bert_large_amp, backend="inductor"), inputs)
 
-    # Run BladeDISC optimization.
-    print("BladeDISC Optimization.")
-    disc_optimize(bert_large_amp, inputs, 'bert_large_amp.disc.pt')
-    model = torch.jit.load('bert_large_amp.disc.pt').cuda().eval()
-    evaluate_torch(model, inputs)
+    # # Run BladeDISC optimization.
+    # print("BladeDISC Optimization.")
+    # disc_optimize(bert_large_amp, inputs, 'bert_large_amp.disc.pt')
+    # model = torch.jit.load('bert_large_amp.disc.pt').cuda().eval()
+    # evaluate_torch(model, inputs)
 
     if args.disc_only:
         return

pytorch_blade/pytorch_blade/torch-mlir/BUILD

@@ -56,7 +56,7 @@ cc_library(
         "lib/Dialect/TorchConversion/Transforms/ReduceTensorConversions.cpp",
         "lib/Dialect/TorchConversion/Transforms/DiscDecomposeComplexOps.cpp",
         "lib/Dialect/TorchConversion/Transforms/DiscConvertTorchToDiscMhlo.cpp",
-        "lib/Dialect/TorchConversion/Transforms/UpgradeLegacyOps.cpp",
+        "lib/Dialect/TorchConversion/Transforms/UpgradeAndAdaptOps.cpp",
         "lib/Dialect/TorchConversion/Transforms/VerifyMhloBackendContract.cpp",
     ],
     hdrs = [

pytorch_blade/pytorch_blade/torch-mlir/include/torch-mlir/Conversion/MhloPasses.h

@@ -70,7 +70,7 @@ std::unique_ptr<OperationPass<func::FuncOp>> createApplyDiscPdlPatternsPass(
 std::unique_ptr<OperationPass<func::FuncOp>> createDiscConvertTorchToMhloPass();
 std::unique_ptr<OperationPass<func::FuncOp>> createDiscDecomposeComplexOpsPass();
 std::unique_ptr<OperationPass<func::FuncOp>> createDiscConvertTorchToDiscMhlo();
-std::unique_ptr<OperationPass<func::FuncOp>> createDiscUpgradeLegacyOpsPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createDiscUpgradeAndAdaptOpsPass();
 std::unique_ptr<OperationPass<ModuleOp>> createVerifyMhloBackendContractPass();
 } // namespace TorchConversion
 } // namespace torch

pytorch_blade/pytorch_blade/torch-mlir/include/torch-mlir/Conversion/MhloPasses.td

@@ -58,15 +58,16 @@ def VerifyMhloBackendContract : Pass<"torch-verify-mhlo-backend-contract", "Modu
   let constructor = "mlir::torch::TorchConversion::createVerifyMhloBackendContractPass()";
 }
 
-def DiscUpgradeLegacyOps : Pass<"torch-disc-upgrade-legacy-ops", "func::FuncOp"> {
-  let summary = "Upgrade Torch legacy operators";
+def DiscUpgradeAndAdaptOps : Pass<"torch-disc-upgrade-and-adapt-ops", "func::FuncOp"> {
+  let summary = "Upgrade Torch legacy operators; Adapt operators as well";
   let description = [{
-    Upgrade outdated torch operations to latest compatiable versions.
+    Upgrade outdated torch operations to latest compatiable versions. Also, used to
+    adapt operations to a correct format.
 
     An example of the transformations done in this pass is:
     - convert aten.gelu(x) to aten.gelu(x, None)
   }];
-  let constructor = "mlir::torch::TorchConversion::createDiscUpgradeLegacyOpsPass()";
+  let constructor = "mlir::torch::TorchConversion::createDiscUpgradeAndAdaptOpsPass()";
 }
 
 def DiscDecomposeComplexOps : Pass<"torch-disc-decompose-complex-ops", "func::FuncOp"> {

pytorch_blade/pytorch_blade/torch-mlir/lib/Conversion/MhloPasses.cpp

@@ -39,7 +39,7 @@ void ::mlir::torch::registerTorchToMhloPasses() {
 void mlir::torch::createDiscTorchBackendToMhloBackendPipeline(
     OpPassManager& pm,
     const Torch::TorchLoweringPipelineOptions& options) {
-  pm.addNestedPass<func::FuncOp>(createDiscUpgradeLegacyOpsPass());
+  pm.addNestedPass<func::FuncOp>(createDiscUpgradeAndAdaptOpsPass());
   ::mlir::torch::Torch::TorchLoweringPipelineOptions funcOptions;
   funcOptions.decompose = false;
   ::mlir::torch::createDiscTorchFunctionToTorchBackendPipeline(pm, funcOptions);

pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/UpgradeLegacyOps.cpp → pytorch_blade/pytorch_blade/torch-mlir/lib/Dialect/TorchConversion/Transforms/UpgradeAndAdaptOps.cpp

@@ -12,6 +12,7 @@
 #include "torch-mlir/Conversion/MhloPasses.h"
 #include "torch-mlir/Dialect/Torch/IR/TorchDialect.h"
 #include "torch-mlir/Dialect/Torch/IR/TorchOps.h"
+#include "torch-mlir/Dialect/Torch/Utils/Utils.h"
 
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Verifier.h"
@@ -62,9 +63,53 @@ LogicalResult UpgradeAtenOp<AtenGeluOp>::matchAndRewrite(
 }
 } // namespace
 
+namespace { // AtenLayerNormOp
+template <>
+bool isLegalAtenOp<AtenLayerNormOp>(AtenLayerNormOp op) {
+  // https://github.com/pytorch/pytorch/issues/66707
+  // layer_norm needs to be done in fp32 for fp16 inputs
+  auto inpDtype =
+      op.input().getType().template dyn_cast<BaseTensorType>().getDtype();
+  if (inpDtype.isa<mlir::FloatType>() &&
+      inpDtype.cast<mlir::FloatType>().getWidth() < 32) {
+    return false;
+  }
+  return true;
+}
+
+template <>
+LogicalResult UpgradeAtenOp<AtenLayerNormOp>::matchAndRewrite(
+    AtenLayerNormOp op,
+    PatternRewriter& rewriter) const {
+  auto inpDtype =
+      op.input().getType().template dyn_cast<BaseTensorType>().getDtype();
+  if (inpDtype.isa<mlir::FloatType>() &&
+      inpDtype.cast<mlir::FloatType>().getWidth() < 32) {
+    Type floatDtype = rewriter.getF32Type();
+    auto loc = op.getLoc();
+    Value input =
+        Torch::convertTensorToDtype(rewriter, loc, op.input(), floatDtype);
+    Value output = rewriter.create<AtenLayerNormOp>(
+        op.getLoc(),
+        input.getType(),
+        input,
+        op.normalized_shape(),
+        op.weight(),
+        op.bias(),
+        op.eps(),
+        op.cudnn_enable());
+    auto outDtype = op.getType().template dyn_cast<BaseTensorType>().getDtype();
+    output = Torch::convertTensorToDtype(rewriter, loc, output, outDtype);
+    rewriter.replaceOp(op, output);
+    return success();
+  }
+  return failure();
+}
+} // namespace
+
 namespace {
-class DiscUpgradeLegacyOpsPass
-    : public DiscUpgradeLegacyOpsBase<DiscUpgradeLegacyOpsPass> {
+class DiscUpgradeAndAdaptOpsPass
+    : public DiscUpgradeAndAdaptOpsBase<DiscUpgradeAndAdaptOpsPass> {
   void runOnOperation() override {
     MLIRContext* context = &getContext();
 
@@ -77,6 +122,7 @@ class DiscUpgradeLegacyOpsPass
   target.addDynamicallyLegalOp<AtenOp>(&isLegalAtenOp<AtenOp>); \
   patterns.add<UpgradeAtenOp<AtenOp>>(context)
     UPGRADE_ATENOP_PATTERN(AtenGeluOp);
+    UPGRADE_ATENOP_PATTERN(AtenLayerNormOp);
 #undef UPGRADE_ATENOP_PATTERN
 
     if (failed(applyPartialConversion(
@@ -89,8 +135,8 @@ class DiscUpgradeLegacyOpsPass
 } //  namespace
 
 std::unique_ptr<OperationPass<func::FuncOp>> mlir::torch::TorchConversion::
-    createDiscUpgradeLegacyOpsPass() {
-  return std::make_unique<DiscUpgradeLegacyOpsPass>();
+    createDiscUpgradeAndAdaptOpsPass() {
+  return std::make_unique<DiscUpgradeAndAdaptOpsPass>();
 }
 
 #undef TORCH_VERSION_LT

pytorch_blade/tests/disc/ops/test_layer_norm.py

@@ -44,5 +44,23 @@ def layernorm(x):
 
         self._test_layer_norm(layernorm)
 
+    def test_layernorm_half_precision(self):
+        @torch.jit.script
+        def layernorm(input, weight, bias):
+            width = 2560
+            eps = 1e-5
+            normalized_shape = (width,)
+            output = torch.layer_norm(input, normalized_shape, weight, bias, eps)
+            return output
+
+        width = 2560
+        input = torch.rand(1, 5, width, device="cuda", dtype=torch.half) * 0.1
+        weight = torch.ones(width, device="cuda", dtype=torch.half)
+        bias = torch.zeros(width, device="cuda", dtype=torch.half)
+
+        annotations = [([-1, -1, -1], torch.half), ([-1], torch.half), ([-1], torch.half)]
+        self._test_disc(layernorm, annotations, (input, weight, bias))
+
+
 if __name__ == "__main__":
     unittest.main()