[llvm-branch-commits] [mlir] 10cf6d4 - [MLIR] Add RngNormalOp in LHLO dialect (#60)

[Uday Bondhugula via llvm-branch-commits]

Author: Abhishek Varma
Date: 2021-09-23T06:20:44+05:30
New Revision: 10cf6d43c473449402675de0d06650464f01bc9f

URL: https://github.com/llvm/llvm-project/commit/10cf6d43c473449402675de0d06650464f01bc9f
DIFF: https://github.com/llvm/llvm-project/commit/10cf6d43c473449402675de0d06650464f01bc9f.diff

LOG: [MLIR] Add RngNormalOp in LHLO dialect (#60)

-- This commit introduces RngNormalOp as part of LHLO dialect.

Signed-off-by: Abhishek Varma <abhishek.varma at polymagelabs.com>

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/LHLO/IR/LHLOOps.td

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/LHLO/IR/LHLOOps.td b/mlir/include/mlir/Dialect/LHLO/IR/LHLOOps.td
index 16fff82d6c7a..394f6c6679a8 100644
--- a/mlir/include/mlir/Dialect/LHLO/IR/LHLOOps.td
+++ b/mlir/include/mlir/Dialect/LHLO/IR/LHLOOps.td
@@ -582,6 +582,26 @@ def LHLO_RngUniformOp : LHLO_Op<"rng_uniform", []>, BASE_HLO_RngUniformOp {
   );
 }
 
+//===----------------------------------------------------------------------===//
+// RngNormal operator that takes scalar low and high values and populates the
+// output memref consistent with HLO op xla_hlo::rng_normal.
+//===----------------------------------------------------------------------===//
+def LHLO_RngNormalOp : LHLO_Op<"rng_normal", []>, BASE_HLO_RngNormalOp {
+  let summary = [{
+   The generated values follow a standard normal distribution in the range 
+   [minval, maxval]. Operations that rely on a random seed actually derive it 
+   from two seeds: the global and operation-level seeds which are `seed` and 
+   `seed2` respectively.
+  }];
+  let arguments = (ins
+    AnyTypeOf<[AnyInteger, AnyFloat]>:$a,
+    AnyTypeOf<[AnyInteger, AnyFloat]>:$b,
+    Arg<LHLO_Buffer, "", [MemWrite]>:$output,
+    OptionalAttr<I64Attr>:$seed,
+    OptionalAttr<I64Attr>:$seed2
+  );
+}
+
 def LHLO_SelectOp: LHLO_Op<"select", []>, BASE_HLO_SelectOp {
   let arguments = (ins
     Arg<LHLO_PredBufferOrI1, "", [MemRead]>:$pred,