EnzymeAD / EnzymeAD/Enzyme-JAX
Test cases for improved for loop raising
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
```mlir
func.func private @"##call__Z29gpu__compute_barotropic_mode_16CompilerMetadataI10StaticSizeI9_180__90_E12DynamicCheckvv7NDRangeILi2ES0_I7_12__6_ES0_I8_16__16_EvvEE5FieldI4Face6Centervvvv11OffsetArrayI7Float64Li3E13CuTracedArrayISC_Li3ELi1E13_194__104__1_EESC_vvvES8_ISA_S9_vvvvSF_SC_vvvE20ImmersedBoundaryGridISC_8Periodic14RightConnected7Bounded28OrthogonalSphericalShellGridISC_SJ_SK_SL_28StaticVerticalDiscretizationISB_ISC_Li1ESD_ISC_Li1ELi1E5_35__EESB_ISC_Li1ESD_ISC_Li1ELi1E5_34__EESP_SR_E8TripolarI5Int64SU_SU_ESB_ISC_Li2ESD_ISC_Li2ELi1E10_194__104_EESX_SX_SX_vSC_E16GridFittedBottomIS8_ISA_SA_vvvvSF_SC_vvvE23CenterImmersedConditionEvvvEvSB_ISC_Li3ESD_ISC_Li3ELi1E14_194__104__34_EES15_SF_#1630$par229"(%arg0: memref<1x104x194xf64, 1>, %arg1: memref<1x104x194xf64, 1>, %arg2: memref<35xf64, 1>, %arg3: memref<34xf64, 1>, %arg4: memref<1x104x194xf64, 1>, %arg5: memref<34x104x194xf64, 1>, %arg6: memref<34x104x194xf64, 1>) {
%cst = arith.constant 0x7FF8000000000000 : f64
%cst_0 = arith.constant 0.000000e+00 : f64
%cst_1 = arith.constant 1.000000e+00 : f64
affine.parallel (%arg7, %arg8) = (0, 0) to (90, 180) {
%0 = affine.load %arg2[27] {alignment = 8 : i64, ordering = 0 : i64, tbaa = [#tbaa_tag]} : memref<35xf64, 1>
%1 = affine.load %arg4[0, %arg7 + 7, %arg8 + 6] : memref<1x104x194xf64, 1>
%2 = arith.subf %0, %1 {fastmathFlags = #llvm.fastmath} : f64
%3 = affine.load %arg4[0, %arg7 + 7, %arg8 + 7] : memref<1x104x194xf64, 1>
%4 = arith.subf %0, %3 {fastmathFlags = #llvm.fastmath} : f64
%5 = math.isnan %2 : f64
%6 = math.isnan %4 : f64
%7 = arith.ori %5, %6 : i1
%8 = arith.minnumf %2, %4 : f64
%9 = arith.select %7, %cst, %8 : f64
%10 = affine.load %arg4[0, %arg7 + 6, %arg8 + 7] : memref<1x104x194xf64, 1>
%11 = arith.subf %0, %10 {fastmathFlags = #llvm.fastmath} : f64
%12 = math.isnan %11 : f64
%13 = arith.ori %12, %6 : i1
%14 = arith.minnumf %11, %4 : f64
%15 = arith.select %13, %cst, %14 : f64
%16 = arith.cmpf oeq, %9, %cst_0 {fastmathFlags = #llvm.fastmath} : f64
%17 = arith.divf %9, %9 {fastmathFlags = #llvm.fastmath} : f64
%18 = arith.select %16, %cst_1, %17 : f64
%19 = arith.cmpf oeq, %15, %cst_0 {fastmathFlags = #llvm.fastmath} : f64
%20 = arith.divf %15, %15 {fastmathFlags = #llvm.fastmath} : f64
%21 = arith.select %19, %cst_1, %20 : f64
%22 = affine.load %arg3[7] {alignment = 8 : i64, ordering = 0 : i64, tbaa = [#tbaa_tag]} : memref<34xf64, 1>
%23 = affine.load %arg5[7, %arg7 + 7, %arg8 + 7] : memref<34x104x194xf64, 1>
%24 = arith.mulf %22, %23 {fastmathFlags = #llvm.fastmath} : f64
%25 = arith.mulf %18, %24 {fastmathFlags = #llvm.fastmath} : f64
affine.store %25, %arg0[0, %arg7 + 7, %arg8 + 7] : memref<1x104x194xf64, 1>
%26 = affine.load %arg3[7] {alignment = 8 : i64, ordering = 0 : i64, tbaa = [#tbaa_tag]} : memref<34xf64, 1>
%27 = affine.load %arg6[7, %arg7 + 7, %arg8 + 7] : memref<34x104x194xf64, 1>
%28 = arith.mulf %26, %27 {fastmathFlags = #llvm.fastmath} : f64
%29 = arith.mulf %21, %28 {fastmathFlags = #llvm.fastmath} : f64
affine.store %29, %arg1[0, %arg7 + 7, %arg8 + 7] : memref<1x104x194xf64, 1>
affine.for %arg9 = 0 to 19 {
%30 = affine.load %arg0[0, %arg7 + 7, %arg8 + 7] : memref<1x104x194xf64, 1>
%31 = affine.load %arg3[%arg9 + 8] {alignment = 8 : i64, ordering = 0 : i64, tbaa = [#tbaa_tag]} : memref<34xf64, 1>
%32 = affine.load %arg5[%arg9 + 8, %arg7 + 7, %arg8 + 7] : memref<34x104x194xf64, 1>
%33 = arith.mulf %31, %32 {fastmathFlags = #llvm.fastmath} : f64
%34 = arith.mulf %18, %33 {fastmathFlags = #llvm.fastmath} : f64
%35 = arith.addf %30, %34 {fastmathFlags = #llvm.fastmath} : f64
affine.store %35, %arg0[0, %arg7 + 7, %arg8 + 7] : memref<1x104x194xf64, 1>
%36 = affine.load %arg1[0, %arg7 + 7, %arg8 + 7] : memref<1x104x194xf64, 1>
%37 = affine.load %arg3[%arg9 + 8] {alignment = 8 : i64, ordering = 0 : i64, tbaa = [#tbaa_tag]} : memref<34xf64, 1>
%38 = affine.load %arg6[%arg9 + 8, %arg7 + 7, %arg8 + 7] : memref<34x104x194xf64, 1>
%39 = arith.mulf %37, %38 {fastmathFlags = #llvm.fastmath} : f64
%40 = arith.mulf %21, %39 {fastmathFlags = #llvm.fastmath} : f64
%41 = arith.addf %36, %40 {fastmathFlags = #llvm.fastmath} : f64
affine.store %41, %arg1[0, %arg7 + 7, %arg8 + 7] : memref<1x104x194xf64, 1>
}
}
return
}
```
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.