[MLIR/GPUToNVVM] lowering arith.maxf/minf to __nv_fmax/__nv_fmin

[qingyunqu]

Summary

Lowering arith.maxf/minf to __nv_fmax/__nv_fmin in GPUToNVVM Conversion.
Also remove Linalg pass declaration which is not cleaned.

[qingyunqu]

Also remove Linalg pass declaration which is not cleaned.

[ThomasRaoux]

mlir/include/mlir/Dialect/Linalg/Passes.h

34

Is this on purpose?

[ThomasRaoux]

mlir/include/mlir/Dialect/Linalg/Passes.h

34

never mind I see your comment and this is dead.

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

257

The semantic of arith::MaxFOp and __nv_fmax_f are different for Nan values:

arith::MaxFOp:

Returns the maximum of the two arguments, treating -0.0 as less than +0.0.
If one of the arguments is NaN, then the result is also NaN.

__nv_fmax_f:

If one argument is a NaN and the other is legitimate numeric value, the numeric
value is chosen.

[qingyunqu]

mlir/include/mlir/Dialect/Linalg/Passes.h

34

this pass seems to be removed at https://reviews.llvm.org/D124145

[qingyunqu]

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

257

This seems to be a problem. BTW, does other math function have the same problem? like math.sin(NaN) and __nv_sin(NaN).

[ThomasRaoux]

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

257

Yes I don't think we can have this lowering with the current semantic of arith::MaxFOp as far as I know this problem is only for fmax/fmin as the semantic of TF and other ML framework tends to be different than what hw natively support. For other arithmetic operations the semantic is more obvious.

[mravishankar]

mlir/include/mlir/Dialect/Linalg/Passes.h

34

Thanks!

[qingyunqu]

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

257

Thanks, should I close this differential?

[ThomasRaoux]

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

257

Yes unfortunately I don’t think there is another solution right now.

mlir/include/mlir/Dialect/Linalg/Passes.h

 struct OneShotBufferizationOptions;
 } // namespace bufferization
 
 std::unique_ptr<Pass> createConvertElementwiseToLinalgPass();
 
 std::unique_ptr<Pass> createLinalgFoldUnitExtentDimsPass();
 
 std::unique_ptr<Pass> createLinalgElementwiseOpFusionPass();
-std::unique_ptr<Pass> createFoldReshapeOpsByLinearizationPass();
 
 std::unique_ptr<Pass> createLinalgNamedOpConversionPass();
 
 std::unique_ptr<OperationPass<func::FuncOp>>
 createLinalgTilingPass(ArrayRef<int64_t> tileSizes = {},
                        linalg::LinalgTilingLoopType loopType =
                            linalg::LinalgTilingLoopType::Loops);
 

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

void mlir::populateGpuToNVVMConversionPatterns(LLVMTypeConverter &converter,
   // Explicitly drop memory space when lowering private memory
   // attributions since NVVM models it as `alloca`s in the default
   // memory space and does not support `alloca`s with addrspace(5).
   patterns.add<GPUFuncOpLowering>(
       converter, /*allocaAddrSpace=*/0,
       StringAttr::get(&converter.getContext(),
                       NVVM::NVVMDialect::getKernelFuncAttrName()));
 
+  patterns.add<OpToFuncCallLowering<arith::MaxFOp>>(converter, "__nv_fmaxf",
+                                                    "__nv_fmax");
+  patterns.add<OpToFuncCallLowering<arith::MinFOp>>(converter, "__nv_fminf",
+                                                    "__nv_fmin");
   patterns.add<OpToFuncCallLowering<math::AbsOp>>(converter, "__nv_fabsf",
                                                   "__nv_fabs");
   patterns.add<OpToFuncCallLowering<math::AtanOp>>(converter, "__nv_atanf",
                                                    "__nv_atan");
   patterns.add<OpToFuncCallLowering<math::Atan2Op>>(converter, "__nv_atan2f",
                                                     "__nv_atan2");
   patterns.add<OpToFuncCallLowering<math::CeilOp>>(converter, "__nv_ceilf",
                                                    "__nv_ceil");

mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir

func.func @gpu_sync() {
     gpu.barrier
     func.return
   }
 }
 
 // -----
 
 gpu.module @test_module {
+  // CHECK: llvm.func @__nv_fmaxf(f32, f32) -> f32
+  // CHECK: llvm.func @__nv_fmax(f64, f64) -> f64
+  // CHECK-LABEL: func @gpu_fmax
+  func.func @gpu_fmax(%arg0_f32 : f32, %arg1_f32 : f32, %arg0_f64 : f64, %arg1_f64 : f64) -> (f32, f64) {
+    %result32 = arith.maxf %arg0_f32, %arg1_f32 : f32
+    // CHECK: llvm.call @__nv_fmaxf(%{{.*}}) : (f32, f32) -> f32
+    %result64 = arith.maxf %arg0_f64, %arg1_f64 : f64
+    // CHECK: llvm.call @__nv_fmax(%{{.*}}) : (f64, f64) -> f64
+    func.return %result32, %result64 : f32, f64
+  }
+}
+
+// -----
+
+gpu.module @test_module {
+  // CHECK: llvm.func @__nv_fminf(f32, f32) -> f32
+  // CHECK: llvm.func @__nv_fmin(f64, f64) -> f64
+  // CHECK-LABEL: func @gpu_fmin
+  func.func @gpu_fmin(%arg0_f32 : f32, %arg1_f32 : f32, %arg0_f64 : f64, %arg1_f64 : f64) -> (f32, f64) {
+    %result32 = arith.minf %arg0_f32, %arg1_f32 : f32
+    // CHECK: llvm.call @__nv_fminf(%{{.*}}) : (f32, f32) -> f32
+    %result64 = arith.minf %arg0_f64, %arg1_f64 : f64
+    // CHECK: llvm.call @__nv_fmin(%{{.*}}) : (f64, f64) -> f64
+    func.return %result32, %result64 : f32, f64
+  }
+}
+
+// -----
+
+gpu.module @test_module {
   // CHECK: llvm.func @__nv_fabsf(f32) -> f32
   // CHECK: llvm.func @__nv_fabs(f64) -> f64
   // CHECK-LABEL: func @gpu_fabs
   func.func @gpu_fabs(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
     %result32 = math.abs %arg_f32 : f32
     // CHECK: llvm.call @__nv_fabsf(%{{.*}}) : (f32) -> f32
     %result64 = math.abs %arg_f64 : f64
     // CHECK: llvm.call @__nv_fabs(%{{.*}}) : (f64) -> f64