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/RefBackend/insert-rng-globals.mlir

21 lines
1.0 KiB
MLIR
Raw Normal View History

[Torch][Linalg] Add basic support for RNG This PR include the following pieces: - Add torch `Generator` type. `Generator` type is converted to i64 in refbackend type converter. - Add seed managment support for the default global generator. `torch_c.getNextSeed` op is used to get the seed. On refbackend, the `torch_c.getNextSeed` is lowered to load/store from [0] of global variable `default_generator` memref<i64> in `InsertRngGlobals` pass. - Add `aten.uniform_` and testing as an example op for RNG ops. Add `torch.pseudo.aten.uniform` op. It has the same operands and return as the `aten.uniform_` from the op registry except for value semantics.
2022-01-12 09:59:42 +08:00
// RUN: torch-mlir-opt %s -refback-insert-rng-globals -split-input-file | FileCheck %s
// CHECK-LABEL: memref.global "private" @global_seed : memref<i64> = dense<0>
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
// CHECK-LABEL: func.func @f() -> i64 {
[Torch][Linalg] Add basic support for RNG This PR include the following pieces: - Add torch `Generator` type. `Generator` type is converted to i64 in refbackend type converter. - Add seed managment support for the default global generator. `torch_c.getNextSeed` op is used to get the seed. On refbackend, the `torch_c.getNextSeed` is lowered to load/store from [0] of global variable `default_generator` memref<i64> in `InsertRngGlobals` pass. - Add `aten.uniform_` and testing as an example op for RNG ops. Add `torch.pseudo.aten.uniform` op. It has the same operands and return as the `aten.uniform_` from the op registry except for value semantics.
2022-01-12 09:59:42 +08:00
// CHECK: %[[MEMREF:.*]] = memref.get_global @global_seed : memref<i64>
// CHECK: %[[SEED:.*]] = memref.load %[[MEMREF]][] : memref<i64>
// CHECK: %[[MULTIPLIER:.*]] = arith.constant 6364136223846793005 : i64
// CHECK: %[[INC:.*]] = arith.constant 1442695040888963407 : i64
// CHECK: %[[MUL:.*]] = arith.muli %[[SEED]], %[[MULTIPLIER]] : i64
// CHECK: %[[TEMP:.*]] = arith.addi %[[MUL]], %[[INC]] : i64
// CHECK: %[[CST127:.*]] = arith.constant 127 : i64
// CHECK: %[[NEXT_SEED:.*]] = arith.andi %[[TEMP]], %[[CST127]] : i64
// CHECK: memref.store %[[NEXT_SEED]], %[[MEMREF]][] : memref<i64>
// CHECK: return %[[NEXT_SEED]] : i64
module {
mlir: bump llvm tag to 5380e3 (#856) In addition to updating the llvm-project submodule, this patch also: 1. updates shape functions and tests so that `func` and `call` operations refer to the `func` dialect 2. avoid duplicate registration of dialects
2022-05-17 03:54:35 +08:00
func.func @f() -> i64 {
[Torch][Linalg] Add basic support for RNG This PR include the following pieces: - Add torch `Generator` type. `Generator` type is converted to i64 in refbackend type converter. - Add seed managment support for the default global generator. `torch_c.getNextSeed` op is used to get the seed. On refbackend, the `torch_c.getNextSeed` is lowered to load/store from [0] of global variable `default_generator` memref<i64> in `InsertRngGlobals` pass. - Add `aten.uniform_` and testing as an example op for RNG ops. Add `torch.pseudo.aten.uniform` op. It has the same operands and return as the `aten.uniform_` from the op registry except for value semantics.
2022-01-12 09:59:42 +08:00
%seed = torch_c.get_next_seed : () -> i64
return %seed : i64
}
}