[Mlir-commits] [mlir] [MLIR][ROCDL] Add dynamically legal ops to LowerGpuOpsToROCDLOpsPass (PR #108302)

[Nirvedh Meshram]

https://github.com/nirvedhmeshram created https://github.com/llvm/llvm-project/pull/108302

Similar to https://github.com/llvm/llvm-project/pull/108266
After https://github.com/llvm/llvm-project/pull/102971
It is legal to generate `LLVM::ExpOp` and `LLVM::LogOp` if the type is f32

>From 4213aa70bdc45abad93c31ceaa28c652ed16b907 Mon Sep 17 00:00:00 2001
From: Nirvedh Meshram <nirvedh at gmail.com>
Date: Wed, 11 Sep 2024 22:04:56 +0000
Subject: [PATCH] [MLIR][ROCDL] Add dynamically legal ops to
 LowerGpuOpsToROCDLOpsPass

---
 mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
index 29926719129dc5..1265f3f48a418e 100644
--- a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
+++ b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
@@ -332,7 +332,14 @@ void mlir::configureGpuToROCDLConversionLegality(ConversionTarget &target) {
   target.addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::Exp2Op, LLVM::FCeilOp,
                       LLVM::FFloorOp, LLVM::FRemOp, LLVM::LogOp, LLVM::Log10Op,
                       LLVM::Log2Op, LLVM::PowOp, LLVM::SinOp>();
-
+  // These ops are legal for f32 type.
+  target.addDynamicallyLegalOp<LLVM::ExpOp, LLVM::LogOp>(
+      [](mlir::Operation *op) {
+        return llvm::any_of(op->getOperandTypes(), [](Type type) {
+          return llvm::isa<FloatType>(type) &&
+                 type.getIntOrFloatBitWidth() == 32;
+        });
+      });
   // TODO: Remove once we support replacing non-root ops.
   target.addLegalOp<gpu::YieldOp, gpu::GPUModuleOp>();
 }