[mlir][rocdl] Add math::Exp2Op lowering to ROCDL

[akuegel]

Summary

There exists already a lowering to LLVM::Exp2Op.
On GPU we prefer to lower to an intrinsic instead.

[herhut]

Thanks!

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

 
 } // anonymous namespace
 
 void mlir::configureGpuToROCDLConversionLegality(ConversionTarget &target) {
   target.addIllegalOp<FuncOp>();
   target.addLegalDialect<::mlir::LLVM::LLVMDialect>();
   target.addLegalDialect<ROCDL::ROCDLDialect>();
   target.addIllegalDialect<gpu::GPUDialect>();
-  target.addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::FAbsOp, LLVM::FCeilOp,
-                      LLVM::FFloorOp, LLVM::LogOp, LLVM::Log10Op, LLVM::Log2Op,
-                      LLVM::PowOp, LLVM::SinOp, LLVM::SqrtOp>();
+  target.addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::Exp2Op, LLVM::FAbsOp,
+                      LLVM::FCeilOp, LLVM::FFloorOp, LLVM::LogOp, LLVM::Log10Op,
+                      LLVM::Log2Op, LLVM::PowOp, LLVM::SinOp, LLVM::SqrtOp>();
 
   // TODO: Remove once we support replacing non-root ops.
   target.addLegalOp<gpu::YieldOp, gpu::GPUModuleOp, gpu::ModuleEndOp>();
 }
 
 void mlir::populateGpuToROCDLConversionPatterns(LLVMTypeConverter &converter,
                                                 RewritePatternSet &patterns) {
   populateWithGenerated(patterns);
void mlir::populateGpuToROCDLConversionPatterns(LLVMTypeConverter &converter,
   patterns.add<OpToFuncCallLowering<math::Atan2Op>>(
       converter, "__ocml_atan2_f32", "__ocml_atan2_f64");
   patterns.add<OpToFuncCallLowering<CeilFOp>>(converter, "__ocml_ceil_f32",
                                               "__ocml_ceil_f64");
   patterns.add<OpToFuncCallLowering<math::CosOp>>(converter, "__ocml_cos_f32",
                                                   "__ocml_cos_f64");
   patterns.add<OpToFuncCallLowering<math::ExpOp>>(converter, "__ocml_exp_f32",
                                                   "__ocml_exp_f64");
+  patterns.add<OpToFuncCallLowering<math::Exp2Op>>(converter, "__ocml_exp2_f32",
+                                                   "__ocml_exp2_f64");
   patterns.add<OpToFuncCallLowering<math::ExpM1Op>>(
       converter, "__ocml_expm1_f32", "__ocml_expm1_f64");
   patterns.add<OpToFuncCallLowering<FloorFOp>>(converter, "__ocml_floor_f32",
                                                "__ocml_floor_f64");
   patterns.add<OpToFuncCallLowering<math::LogOp>>(converter, "__ocml_log_f32",
                                                   "__ocml_log_f64");
   patterns.add<OpToFuncCallLowering<math::Log10Op>>(
       converter, "__ocml_log10_f32", "__ocml_log10_f64");

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

func @gpu_exp(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
     %result32 = math.exp %exp_f32 : f32
     // CHECK: llvm.call @__ocml_exp_f32(%{{.*}}) : (f32) -> f32
     %result64 = math.exp %arg_f64 : f64
     // CHECK: llvm.call @__ocml_exp_f64(%{{.*}}) : (f64) -> f64
     std.return %result32, %result64 : f32, f64
   }
 }
 
+// -----
+
+gpu.module @test_module {
+  // CHECK: llvm.func @__ocml_exp2_f32(f32) -> f32
+  // CHECK: llvm.func @__ocml_exp2_f64(f64) -> f64
+  // CHECK-LABEL: func @gpu_exp2
+  func @gpu_exp2(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
+    %exp2_f32 = math.exp2 %arg_f32 : f32
+    // CHECK: llvm.call @__ocml_exp2_f32(%{{.*}}) : (f32) -> f32
+    %result32 = math.exp2 %exp2_f32 : f32
+    // CHECK: llvm.call @__ocml_exp2_f32(%{{.*}}) : (f32) -> f32
+    %result64 = math.exp2 %arg_f64 : f64
+    // CHECK: llvm.call @__ocml_exp2_f64(%{{.*}}) : (f64) -> f64
+    std.return %result32, %result64 : f32, f64
+  }
+}
 
 // -----
 
 // 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 @__ocml_exp_f32(f32) -> f32