[llvm-branch-commits] [mlir] db04767 - [MLIR] Add the seed attribute for random_uniform ops in lmhlo dialect. (#35)

[Uday Bondhugula via llvm-branch-commits]

Author: Prashant Kumar
Date: 2021-09-23T06:20:43+05:30
New Revision: db04767afbea48b62a7a25017e42248a5d83480c

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

LOG: [MLIR] Add the seed attribute for random_uniform ops in lmhlo dialect. (#35)

The `seed` and `seed2` attribute are added to lmhlo dialect in case of
random_uniform ops. The changes are made as a part of hlo-legalize-to-lhlo` pass.

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 0ee24446cad8..16fff82d6c7a 100644
--- a/mlir/include/mlir/Dialect/LHLO/IR/LHLOOps.td
+++ b/mlir/include/mlir/Dialect/LHLO/IR/LHLOOps.td
@@ -568,10 +568,17 @@ def LHLO_ReshapeOp: LHLO_Op<"reshape", []>, BASE_HLO_ReshapeOp {
 // output memref consistent with HLO op xla_hlo::rng_uniform.
 //===----------------------------------------------------------------------===//
 def LHLO_RngUniformOp : LHLO_Op<"rng_uniform", []>, BASE_HLO_RngUniformOp {
+  let summary = [{
+   The generated values follow a uniform 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
+    Arg<LHLO_Buffer, "", [MemWrite]>:$output,
+    OptionalAttr<I64Attr>:$seed,
+    OptionalAttr<I64Attr>:$seed2
   );
 }