Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add mm, sqrt, div operators for complex tensor to support OneScience #10269

Merged
merged 219 commits into from
Jun 30, 2023

Conversation

MarioLulab
Copy link
Contributor

@MarioLulab MarioLulab commented May 16, 2023

Original requirements

AI for Science models FNO use a series of operators such as sqrt, mm, div, etc. for complex tensor. So we need to support complex tensor for these ops.

Main Works

Support existing ops for complex tensor in this pr as below:

  • broadcast_elementwise_binary

    Op complex type Backend Progress
    scalar_div cp64, cp128 CPU, CUDA Done
    broadcast_div cp64, cp128 CPU, CUDA Done
  • broadcast_elementwise_unary

    Op complex type Backend Progress
    sqrt cp64, cp128 CPU, CUDA Done
    negative cp64, cp128 CPU, CUDA Done
  • other exisiting operations

    Op complex type Backend Progress
    matmul cp64, cp128 CPU, CUDA Done
    batch_matmul cp64, cp128 CPU, CUDA Done
    broadcast_matmul cp64, cp128 CPU, CUDA Done
    matrix_vector_product cp64, cp128 CPU, CUDA Done
    reduce_sum_like cp64, cp128 CPU, CUDA Done

注意:

复数基础设施建设系列 pr:

  1. 使用 primitive 来实现 conj, real 等常见复数算子: use primitive to implement "conj_physical", "real" and "imag" op #10281
  2. 将现有支持复数数据类型的算子测例迁移到 autotest 模块中,以使复数算子复用实数算子的测试用例:Applying autotest module on existing complex operators #10284
  3. 继续拓展支持复数数据类型的算子,比如 matmul, sqrt, div 等:add mm, sqrt, div operators for complex tensor to support OneScience #10269
    依赖关系:
    本 pr 基于:pr1pr2 ,需要在 merge pr1pr2 后,再 Merge 本 pr

levi131 pushed a commit that referenced this pull request Jun 5, 2023
)

### Original requirements
现有的 conj kernel, real kernel, imag kernel 使用 KernelUtil 来进行组织,在 kernel
层面复数数据类型和实数数据类型的调用 conj 无法实现统一。故移除原有 KernelUtil 的实现,使用 ElementwiseUnary
的 primitive 进行实现

- [x] 新增 `kConj` 这一 UnaryOp,使用 ElementwiseUnary 的 Primitive 实现
conj,并移除原有 conj 的 KernelUtil
- [x] 新增 `kReal` 和 `kRealGrad` UnaryOp,使用 ElementwiseUnary 的 Primitive
实现 real 和 real_grad,并移除原有 real 和 real_grad 的 KernelUtil
- [x] 新增 `kImag` 和 `kImagGrad` UnaryOp,使用 ElementwiseUnary 的 Primitive
实现 imag 和 imag_grad,并移除原有 imag 和 imag_grad 的 KernelUtil

## 注意:
复数基础设施建设系列 pr:
1. 使用 primitive 来实现 conj, real 等常见复数算子:
#10281
2. 将现有支持复数数据类型的算子测例迁移到 autotest
模块中,以使复数算子复用实数算子的测试用例:#10284
3. 继续拓展支持复数数据类型的算子,比如 matmul, sqrt, div
等:#10269
依赖关系:
本 pr 基于:[pr2](#10284) 和
[pr3](#10269) 的基础,请优先 merge 此
pr
L-Xiafeng pushed a commit that referenced this pull request Jun 26, 2023
### Original requirements
**Autotest**: We found that the previous testing of operators supporting
complex tensor was not complete. We decided to reuse the real tensor
operator tests to ensure completeness. Since complex tensor tests are
supported in the `autotest` module from pr
(#10027) , in this pr we
applied the autotest module to the tests of complex tensor operators
already available in Oneflow

**Fix**: In addition, the autograd rules for some previous operators of
complex numbers might not conform to the convention of ["Conjugate
Wirtinger
Derivative"](https://en.wikipedia.org/wiki/Wirtinger_derivatives). We
have fixed these bugs in this pr at the same time.

#### Main Works
**Applying `autotest` module on existing operators that have already
support complex tensor:**

`Complex and Real Behave the Same Way`: means we don't need to add
conjugate operation in op grad. Because regardless of whether the input
data type involved in the operation is real or complex, the gradient
result using the winterger derivative is the same as the real derivative
rule,

`Grad Not Supported in OF`: means the grad of this op is not supported
in oneflow

- broadcast_elementwise_binary
| Op | complex type | Backend | Using autotest | conjugate Wirtinger
derivative |
  |:-----:|:-------------:|:------:|:------:|:------:|
| Add | cp64, cp128 | CPU, CUDA | DONE | Complex and Real Behave the
Same Way |
  | Mul |  cp64, cp128   |  CPU, CUDA | DONE | DONE |
| Sub | cp64, cp128 | CPU, CUDA | DONE | Complex and Real Behave the
Same Way |
| Equal | cp64, cp128 | CPU, CUDA | DONE | Complex and Real Behave the
Same Way |
| NotEqual | cp64, cp128 | CPU, CUDA | DONE | Grad Not Supported in OF |


- broadcast_elementwise_unary
| Op | complex type | Backend | Using autotest | conjugate Wirtinger
derivative |
  |:----------:|:-------------:|:------:|:------:|:------:|
| Cast | cp64, cp128 | CPU, CUDA | DONE | Complex and Real Behave the
Same Way |


- other exisiting operations
| Op | complex type | Backend | Using autotest | conjugate Wirtinger
derivative |
  |:----------:|:-------------:|:------:|:------:|:------:|
| constant_pad | cp64, cp128 | CPU, CUDA | Done | Complex and Real
Behave the Same Way |
| reduce_sum | cp64, cp128 | CPU, CUDA | TO-DO | Complex and Real Behave
the Same Way |


## 注意:
复数基础设施建设系列 pr:
1. 使用 primitive 来实现 conj, real 等常见复数算子:
#10281
2. 将现有支持复数数据类型的算子测例迁移到 autotest
模块中,以使复数算子复用实数算子的测试用例:#10284
3. 继续拓展支持复数数据类型的算子,比如 matmul, sqrt, div
等:#10269
依赖关系:
本 pr 基于:[pr1](https://github.com/Oneflow-Inc/oneflow/pull/10281),需要在
merge [pr1](#10281) 后,再 Merge
本 pr
@L-Xiafeng L-Xiafeng requested review from oneflow-ci-bot and removed request for oneflow-ci-bot June 28, 2023 06:11
set test_complex
complex test
test complex
test complex
skip test out of memory
@L-Xiafeng L-Xiafeng enabled auto-merge (squash) June 30, 2023 09:30
Copy link
Contributor

@levi131 levi131 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@L-Xiafeng L-Xiafeng merged commit 11d6547 into Oneflow-Inc:master Jun 30, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants