diff --git a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp --- a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp +++ b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp @@ -712,7 +712,8 @@ populateGpuToNVVMConversionPatterns(converter, patterns); ConversionTarget target(getContext()); target.addIllegalDialect(); - target.addIllegalOp(); + target.addIllegalOp(); target.addIllegalOp(); target.addLegalDialect(); target.addLegalDialect(); @@ -739,6 +740,12 @@ NVVM::GridDimYOp, NVVM::GridDimZOp>, GPUAllReduceOpLowering, GPUShuffleOpLowering, GPUFuncOpLowering, GPUReturnOpLowering>(converter); + patterns.insert>(converter, "__nv_fabsf", + "__nv_fabs"); + patterns.insert>(converter, "__nv_ceilf", + "__nv_ceil"); + patterns.insert>(converter, "__nv_cosf", + "__nv_cos"); patterns.insert>(converter, "__nv_expf", "__nv_exp"); } diff --git a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp --- a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp +++ b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp @@ -51,12 +51,19 @@ GPUIndexIntrinsicOpLowering>( converter); + patterns.insert>(converter, "_ocml_fabs_f32", + "_ocml_fabs_f64"); + patterns.insert>(converter, "_ocml_ceil_f32", + "_ocml_ceil_f64"); + patterns.insert>(converter, "_ocml_cos_f32", + "_ocml_cos_f64"); patterns.insert>(converter, "_ocml_exp_f32", "_ocml_exp_f64"); ConversionTarget target(getContext()); target.addLegalDialect(); - target.addIllegalOp(); + target.addIllegalOp(); target.addDynamicallyLegalOp( [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); }); if (failed(applyPartialConversion(m, target, patterns, &converter))) diff --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir --- a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir +++ b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir @@ -112,6 +112,51 @@ // ----- module attributes {gpu.kernel_module} { + // CHECK: llvm.func @__nv_fabsf(!llvm.float) -> !llvm.float + // CHECK: llvm.func @__nv_fabs(!llvm.double) -> !llvm.double + // CHECK-LABEL: func @gpu_fabs + func @gpu_fabs(%arg_f32 : f32, %arg_f64 : f64) { + %result32 = std.absf %arg_f32 : f32 + // CHECK: llvm.call @__nv_fabsf(%{{.*}}) : (!llvm.float) -> !llvm.float + %result64 = std.absf %arg_f64 : f64 + // CHECK: llvm.call @__nv_fabs(%{{.*}}) : (!llvm.double) -> !llvm.double + std.return + } +} + +// ----- + +module attributes {gpu.kernel_module} { + // CHECK: llvm.func @__nv_ceilf(!llvm.float) -> !llvm.float + // CHECK: llvm.func @__nv_ceil(!llvm.double) -> !llvm.double + // CHECK-LABEL: func @gpu_ceil + func @gpu_ceil(%arg_f32 : f32, %arg_f64 : f64) { + %result32 = std.ceilf %arg_f32 : f32 + // CHECK: llvm.call @__nv_ceilf(%{{.*}}) : (!llvm.float) -> !llvm.float + %result64 = std.ceilf %arg_f64 : f64 + // CHECK: llvm.call @__nv_ceil(%{{.*}}) : (!llvm.double) -> !llvm.double + std.return + } +} + +// ----- + +module attributes {gpu.kernel_module} { + // CHECK: llvm.func @__nv_cosf(!llvm.float) -> !llvm.float + // CHECK: llvm.func @__nv_cos(!llvm.double) -> !llvm.double + // CHECK-LABEL: func @gpu_cos + func @gpu_cos(%arg_f32 : f32, %arg_f64 : f64) { + %result32 = std.cos %arg_f32 : f32 + // CHECK: llvm.call @__nv_cosf(%{{.*}}) : (!llvm.float) -> !llvm.float + %result64 = std.cos %arg_f64 : f64 + // CHECK: llvm.call @__nv_cos(%{{.*}}) : (!llvm.double) -> !llvm.double + std.return + } +} + +// ----- + +module attributes {gpu.kernel_module} { // CHECK: llvm.func @__nv_expf(!llvm.float) -> !llvm.float // CHECK: llvm.func @__nv_exp(!llvm.double) -> !llvm.double // CHECK-LABEL: func @gpu_exp diff --git a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir --- a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir +++ b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir @@ -39,6 +39,51 @@ // ----- module attributes {gpu.kernel_module} { + // CHECK: llvm.func @_ocml_fabs_f32(!llvm.float) -> !llvm.float + // CHECK: llvm.func @_ocml_fabs_f64(!llvm.double) -> !llvm.double + // CHECK-LABEL: func @gpu_fabs + func @gpu_fabs(%arg_f32 : f32, %arg_f64 : f64) { + %result32 = std.absf %arg_f32 : f32 + // CHECK: llvm.call @_ocml_fabs_f32(%{{.*}}) : (!llvm.float) -> !llvm.float + %result64 = std.absf %arg_f64 : f64 + // CHECK: llvm.call @_ocml_fabs_f64(%{{.*}}) : (!llvm.double) -> !llvm.double + std.return + } +} + +// ----- + +module attributes {gpu.kernel_module} { + // CHECK: llvm.func @_ocml_ceil_f32(!llvm.float) -> !llvm.float + // CHECK: llvm.func @_ocml_ceil_f64(!llvm.double) -> !llvm.double + // CHECK-LABEL: func @gpu_ceil + func @gpu_ceil(%arg_f32 : f32, %arg_f64 : f64) { + %result32 = std.ceilf %arg_f32 : f32 + // CHECK: llvm.call @_ocml_ceil_f32(%{{.*}}) : (!llvm.float) -> !llvm.float + %result64 = std.ceilf %arg_f64 : f64 + // CHECK: llvm.call @_ocml_ceil_f64(%{{.*}}) : (!llvm.double) -> !llvm.double + std.return + } +} + +// ----- + +module attributes {gpu.kernel_module} { + // CHECK: llvm.func @_ocml_cos_f32(!llvm.float) -> !llvm.float + // CHECK: llvm.func @_ocml_cos_f64(!llvm.double) -> !llvm.double + // CHECK-LABEL: func @gpu_cos + func @gpu_cos(%arg_f32 : f32, %arg_f64 : f64) { + %result32 = std.cos %arg_f32 : f32 + // CHECK: llvm.call @_ocml_cos_f32(%{{.*}}) : (!llvm.float) -> !llvm.float + %result64 = std.cos %arg_f64 : f64 + // CHECK: llvm.call @_ocml_cos_f64(%{{.*}}) : (!llvm.double) -> !llvm.double + std.return + } +} + +// ----- + +module attributes {gpu.kernel_module} { // CHECK: llvm.func @_ocml_exp_f32(!llvm.float) -> !llvm.float // CHECK: llvm.func @_ocml_exp_f64(!llvm.double) -> !llvm.double // CHECK-LABEL: func @gpu_exp