Add GPU lowerings for the different log ops.

[akuegel]

Summary

This adds GPU lowerings for log, log10 and log2.

[akuegel]

Landed as 39e1c1fa9ee03e91751e505d747275e58069e6de

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

converter.addConversion([&](MemRefType type) -> Optional<Type> {
       return converter.convertType(MemRefType::Builder(type).setMemorySpace(0));
     });
 
     OwningRewritePatternList patterns;
     populateStdToLLVMConversionPatterns(converter, patterns);
     populateGpuToNVVMConversionPatterns(converter, patterns);
     ConversionTarget target(getContext());
     target.addIllegalDialect<gpu::GPUDialect>();
-    target.addIllegalOp<LLVM::FAbsOp, LLVM::FCeilOp, LLVM::CosOp,
-                        LLVM::ExpOp>();
+    target.addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::FAbsOp, LLVM::FCeilOp,
+                        LLVM::LogOp, LLVM::Log10Op, LLVM::Log2Op>();
     target.addIllegalOp<FuncOp>();
     target.addLegalDialect<LLVM::LLVMDialect>();
     target.addLegalDialect<NVVM::NVVMDialect>();
     target.addDynamicallyLegalOp<mlir::LLVM::CallOp>(
         gpu::filterIllegalLLVMIntrinsics({"tanh", "tanhf"}, m.getContext()));
     // TODO(csigg): Remove once we support replacing non-root ops.
     target.addLegalOp<gpu::YieldOp, gpu::GPUModuleOp, gpu::ModuleEndOp>();
     if (failed(applyPartialConversion(m, target, patterns, &converter)))
void mlir::populateGpuToNVVMConversionPatterns(
   patterns.insert<OpToFuncCallLowering<AbsFOp>>(converter, "__nv_fabsf",
                                                "__nv_fabs");
   patterns.insert<OpToFuncCallLowering<CeilFOp>>(converter, "__nv_ceilf",
                                                "__nv_ceil");
   patterns.insert<OpToFuncCallLowering<CosOp>>(converter, "__nv_cosf",
                                                "__nv_cos");
   patterns.insert<OpToFuncCallLowering<ExpOp>>(converter, "__nv_expf",
                                                "__nv_exp");
+  patterns.insert<OpToFuncCallLowering<LogOp>>(converter, "__nv_logf",
+                                               "__nv_log");
+  patterns.insert<OpToFuncCallLowering<Log10Op>>(converter, "__nv_log10f",
+                                                 "__nv_log10");
+  patterns.insert<OpToFuncCallLowering<Log2Op>>(converter, "__nv_log2f",
+                                                "__nv_log2");
   patterns.insert<OpToFuncCallLowering<TanhOp>>(converter, "__nv_tanhf",
                                                 "__nv_tanh");
 }
 
 std::unique_ptr<OpPassBase<gpu::GPUModuleOp>>
 mlir::createLowerGpuOpsToNVVMOpsPass() {
   return std::make_unique<LowerGpuOpsToNVVMOpsPass>();
 }
 
 static PassRegistration<LowerGpuOpsToNVVMOpsPass>
     pass("convert-gpu-to-nvvm", "Generate NVVM operations for gpu operations");

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

void runOnOperation() override {
     patterns.insert<OpToFuncCallLowering<AbsFOp>>(converter, "__ocml_fabs_f32",
                                                   "__ocml_fabs_f64");
     patterns.insert<OpToFuncCallLowering<CeilFOp>>(converter, "__ocml_ceil_f32",
                                                    "__ocml_ceil_f64");
     patterns.insert<OpToFuncCallLowering<CosOp>>(converter, "__ocml_cos_f32",
                                                  "__ocml_cos_f64");
     patterns.insert<OpToFuncCallLowering<ExpOp>>(converter, "__ocml_exp_f32",
                                                  "__ocml_exp_f64");
+    patterns.insert<OpToFuncCallLowering<LogOp>>(converter, "__ocml_log_f32",
+                                                 "__ocml_log_f64");
+    patterns.insert<OpToFuncCallLowering<Log10Op>>(
+        converter, "__ocml_log10_f32", "__ocml_log10_f64");
+    patterns.insert<OpToFuncCallLowering<Log2Op>>(converter, "__ocml_log2_f32",
+                                                  "__ocml_log2_f64");
     patterns.insert<OpToFuncCallLowering<TanhOp>>(converter, "__ocml_tanh_f32",
                                                   "__ocml_tanh_f64");
 
     ConversionTarget target(getContext());
     target.addLegalDialect<LLVM::LLVMDialect, ROCDL::ROCDLDialect>();
-    target.addIllegalOp<LLVM::FAbsOp, LLVM::FCeilOp, LLVM::CosOp,
-                        LLVM::ExpOp>();
+    target.addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::FAbsOp, LLVM::FCeilOp,
+                        LLVM::LogOp, LLVM::Log10Op, LLVM::Log2Op>();
     target.addDynamicallyLegalOp<LLVM::CallOp>(
         gpu::filterIllegalLLVMIntrinsics({"tanh", "tanhf"}, m.getContext()));
     target.addIllegalOp<FuncOp>();
     if (failed(applyPartialConversion(m, target, patterns, &converter)))
       signalPassFailure();
   }
 };
 

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

func @gpu_cos(%arg_f32 : f32, %arg_f64 : f64) {
     // 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
   }
 }
 
 // -----
-
-gpu.module @test_module {
-  // CHECK: llvm.func @__nv_tanhf(!llvm.float) -> !llvm.float
-  // CHECK: llvm.func @__nv_tanh(!llvm.double) -> !llvm.double
-  // CHECK-LABEL: func @gpu_tanh
-  func @gpu_tanh(%arg_f32 : f32, %arg_f64 : f64) {
-    %result32 = std.tanh %arg_f32 : f32
-    // CHECK: llvm.call @__nv_tanhf(%{{.*}}) : (!llvm.float) -> !llvm.float
-    %result64 = std.tanh %arg_f64 : f64
-    // CHECK: llvm.call @__nv_tanh(%{{.*}}) : (!llvm.double) -> !llvm.double
-    std.return
-  }
-}
-
-// -----
 gpu.module @test_module {
   // CHECK: llvm.func @__nv_expf(!llvm.float) -> !llvm.float
   // CHECK: llvm.func @__nv_exp(!llvm.double) -> !llvm.double
   // CHECK-LABEL: func @gpu_exp
   func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) {
     %exp_f32 = std.exp %arg_f32 : f32
     // CHECK: llvm.call @__nv_expf(%{{.*}}) : (!llvm.float) -> !llvm.float
     %result_f32 = std.exp %exp_f32 : f32
     // CHECK: llvm.call @__nv_expf(%{{.*}}) : (!llvm.float) -> !llvm.float
     %result64 = std.exp %arg_f64 : f64
     // CHECK: llvm.call @__nv_exp(%{{.*}}) : (!llvm.double) -> !llvm.double
     std.return
   }
 }
 
 // -----
 
+gpu.module @test_module {
+  // CHECK: llvm.func @__nv_logf(!llvm.float) -> !llvm.float
+  // CHECK: llvm.func @__nv_log(!llvm.double) -> !llvm.double
+  // CHECK-LABEL: func @gpu_log
+  func @gpu_log(%arg_f32 : f32, %arg_f64 : f64) {
+    %result32 = std.log %arg_f32 : f32
+    // CHECK: llvm.call @__nv_logf(%{{.*}}) : (!llvm.float) -> !llvm.float
+    %result64 = std.log %arg_f64 : f64
+    // CHECK: llvm.call @__nv_log(%{{.*}}) : (!llvm.double) -> !llvm.double
+    std.return
+  }
+}
+
+// -----
+
+gpu.module @test_module {
+  // CHECK: llvm.func @__nv_log10f(!llvm.float) -> !llvm.float
+  // CHECK: llvm.func @__nv_log10(!llvm.double) -> !llvm.double
+  // CHECK-LABEL: func @gpu_log10
+  func @gpu_log10(%arg_f32 : f32, %arg_f64 : f64) {
+    %result32 = std.log10 %arg_f32 : f32
+    // CHECK: llvm.call @__nv_log10f(%{{.*}}) : (!llvm.float) -> !llvm.float
+    %result64 = std.log10 %arg_f64 : f64
+    // CHECK: llvm.call @__nv_log10(%{{.*}}) : (!llvm.double) -> !llvm.double
+    std.return
+  }
+}
+
+// -----
+
+gpu.module @test_module {
+  // CHECK: llvm.func @__nv_log2f(!llvm.float) -> !llvm.float
+  // CHECK: llvm.func @__nv_log2(!llvm.double) -> !llvm.double
+  // CHECK-LABEL: func @gpu_log2
+  func @gpu_log2(%arg_f32 : f32, %arg_f64 : f64) {
+    %result32 = std.log2 %arg_f32 : f32
+    // CHECK: llvm.call @__nv_log2f(%{{.*}}) : (!llvm.float) -> !llvm.float
+    %result64 = std.log2 %arg_f64 : f64
+    // CHECK: llvm.call @__nv_log2(%{{.*}}) : (!llvm.double) -> !llvm.double
+    std.return
+  }
+}
+
+// -----
+
+gpu.module @test_module {
+  // CHECK: llvm.func @__nv_tanhf(!llvm.float) -> !llvm.float
+  // CHECK: llvm.func @__nv_tanh(!llvm.double) -> !llvm.double
+  // CHECK-LABEL: func @gpu_tanh
+  func @gpu_tanh(%arg_f32 : f32, %arg_f64 : f64) {
+    %result32 = std.tanh %arg_f32 : f32
+    // CHECK: llvm.call @__nv_tanhf(%{{.*}}) : (!llvm.float) -> !llvm.float
+    %result64 = std.tanh %arg_f64 : f64
+    // CHECK: llvm.call @__nv_tanh(%{{.*}}) : (!llvm.double) -> !llvm.double
+    std.return
+  }
+}
+
+// -----
+
 // Test that we handled properly operation with SymbolTable other than module op
 gpu.module @test_module {
   "test.symbol_scope"() ({
   // CHECK: test.symbol_scope
   // CHECK: llvm.func @__nv_expf(!llvm.float) -> !llvm.float
   // CHECK: llvm.func @__nv_exp(!llvm.double) -> !llvm.double
   // CHECK-LABEL: func @gpu_exp
     func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) {

mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir

func @gpu_cos(%arg_f32 : f32, %arg_f64 : f64) {
     // 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
   }
 }
 
 // -----
-
-gpu.module @kernel_module {
-  // CHECK: llvm.func @__ocml_tanh_f32(!llvm.float) -> !llvm.float
-  // CHECK: llvm.func @__ocml_tanh_f64(!llvm.double) -> !llvm.double
-  // CHECK-LABEL: func @gpu_tanh
-  func @gpu_tanh(%arg_f32 : f32, %arg_f64 : f64) {
-    %result32 = std.tanh %arg_f32 : f32
-    // CHECK: llvm.call @__ocml_tanh_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
-    %result64 = std.tanh %arg_f64 : f64
-    // CHECK: llvm.call @__ocml_tanh_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
-    std.return
-  }
-}
-
-// -----
 gpu.module @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
   func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) {
     %exp_f32 = std.exp %arg_f32 : f32
     // CHECK: llvm.call @__ocml_exp_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
     %result_f32 = std.exp %exp_f32 : f32
func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) {
       // CHECK: llvm.call @__ocml_exp_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
       %result64 = std.exp %arg_f64 : f64
       // CHECK: llvm.call @__ocml_exp_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
       std.return
     }
     "test.finish" () : () -> ()
   }) : () -> ()
 }
+
+// -----
+
+gpu.module @kernel_module {
+  // CHECK: llvm.func @__ocml_log_f32(!llvm.float) -> !llvm.float
+  // CHECK: llvm.func @__ocml_log_f64(!llvm.double) -> !llvm.double
+  // CHECK-LABEL: func @gpu_log
+  func @gpu_log(%arg_f32 : f32, %arg_f64 : f64) {
+    %result32 = std.log %arg_f32 : f32
+    // CHECK: llvm.call @__ocml_log_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
+    %result64 = std.log %arg_f64 : f64
+    // CHECK: llvm.call @__ocml_log_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
+    std.return
+  }
+}
+
+// -----
+
+gpu.module @kernel_module {
+  // CHECK: llvm.func @__ocml_log10_f32(!llvm.float) -> !llvm.float
+  // CHECK: llvm.func @__ocml_log10_f64(!llvm.double) -> !llvm.double
+  // CHECK-LABEL: func @gpu_log10
+  func @gpu_log10(%arg_f32 : f32, %arg_f64 : f64) {
+    %result32 = std.log10 %arg_f32 : f32
+    // CHECK: llvm.call @__ocml_log10_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
+    %result64 = std.log10 %arg_f64 : f64
+    // CHECK: llvm.call @__ocml_log10_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
+    std.return
+  }
+}
+
+// -----
+
+gpu.module @kernel_module {
+  // CHECK: llvm.func @__ocml_log2_f32(!llvm.float) -> !llvm.float
+  // CHECK: llvm.func @__ocml_log2_f64(!llvm.double) -> !llvm.double
+  // CHECK-LABEL: func @gpu_log2
+  func @gpu_log2(%arg_f32 : f32, %arg_f64 : f64) {
+    %result32 = std.log2 %arg_f32 : f32
+    // CHECK: llvm.call @__ocml_log2_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
+    %result64 = std.log2 %arg_f64 : f64
+    // CHECK: llvm.call @__ocml_log2_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
+    std.return
+  }
+}
+
+// -----
+
+gpu.module @kernel_module {
+  // CHECK: llvm.func @__ocml_tanh_f32(!llvm.float) -> !llvm.float
+  // CHECK: llvm.func @__ocml_tanh_f64(!llvm.double) -> !llvm.double
+  // CHECK-LABEL: func @gpu_tanh
+  func @gpu_tanh(%arg_f32 : f32, %arg_f64 : f64) {
+    %result32 = std.tanh %arg_f32 : f32
+    // CHECK: llvm.call @__ocml_tanh_f32(%{{.*}}) : (!llvm.float) -> !llvm.float
+    %result64 = std.tanh %arg_f64 : f64
+    // CHECK: llvm.call @__ocml_tanh_f64(%{{.*}}) : (!llvm.double) -> !llvm.double
+    std.return
+  }
+}