[MLIR][GPU] Disallow llvm tanh intrinsics when lowering to NVVM/ROCm.

[herhut]

Summary

The lowering to NVVM and ROCm handles tanh operations differently by
mapping them to NVVM/ROCm specific intrinsics. This conflicts with
the lowering to LLVM, which uses the default llvm intrinsic. This change
declares the LLVM intrinsics to be illegal, hence disallowing the
correspondign rewrite.

[herhut]

Explicit lambda instead of std::bind.

[herhut]

PTAL.

Writing this in code for discussion. Is this the way to resolve conflicts or do we want to hand-pick patterns from the LLVM conversion? The latter seems brittle as well, as the number of pattern keeps growing and we need to track. By disallowing, we can do this whenever we add a new intrinsic.

I am somewhat torn between the two myself, so I wrote it down to see what it looks like. This would unbreak the windows builds.

[ftynse]

This would work for me. At scale, we may need to analyze effects of such approach on compile time. The converter essentially twice the number of patterns it has to search through, and there are potential rollbacks of rewrites that generated illegal operations.

In general, I think we want to expose individual LLVM patterns anyway because several people hit the selection problem repeatedly. I haven't had time to do that yet.

For intrinsics specifically, I have a vague idea of some fallback mechanism, as in: the default Std->LLVM-intrinsic patterns are enabled unless there is a different pattern.

mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h

103

Nit: SmallVector is re-exported into the mlir namespace, drop the llvm::

[herhut]

Drop llvm.

[herhut]

In D74389#1868977, @ftynse wrote:

This would work for me. At scale, we may need to analyze effects of such approach on compile time. The converter essentially twice the number of patterns it has to search through, and there are potential rollbacks of rewrites that generated illegal operations.

In general, I think we want to expose individual LLVM patterns anyway because several people hit the selection problem repeatedly. I haven't had time to do that yet.

The granularity of exposing them is not clear to me. I also thought about an intrinsics/non-intrinsics split. But even the definition of what intrinsic based lowering means is vague. LLVM::ExpOp is an operation but ultimately lowered to some intrinsic. One way to make this cleaner would be to expose all intrinsics as LLVM::Operation in MLIR and only lower them when going to llvm proper. Then at least we can filter on operations and not intrinsic names.

For intrinsics specifically, I have a vague idea of some fallback mechanism, as in: the default Std->LLVM-intrinsic patterns are enabled unless there is a different pattern.

It would be great if one could define priorities of patterns so that certain pattern shadow others. This is different from benefit in that it would have stronger semantics (disallowing patterns rather than preferring them). One idea would be to group pattern into formal groups and then allowing to specify a partial order on these groups, which is then used to specify application order of patterns. That would solve out case where NVVM pattern are meant to shadow LLVM pattern.

For now, should we land this?

[ftynse]

In D74389#1869018, @herhut wrote:

The granularity of exposing them is not clear to me. I also thought about an intrinsics/non-intrinsics split. But even the definition of what intrinsic based lowering means is vague. LLVM::ExpOp is an operation but ultimately lowered to some intrinsic.

The notion of intrinsic does not exist in MLIR proper, everything is an operation. I replicated it in the LLVM dialect to better match the LLVM IR, and to commonalize the lowering.

One way to make this cleaner would be to expose all intrinsics as LLVM::Operation in MLIR and only lower them when going to llvm proper. Then at least we can filter on operations and not intrinsic names.

We should absolutely do that.

For intrinsics specifically, I have a vague idea of some fallback mechanism, as in: the default Std->LLVM-intrinsic patterns are enabled unless there is a different pattern.

It would be great if one could define priorities of patterns so that certain pattern shadow others. This is different from benefit in that it would have stronger semantics (disallowing patterns rather than preferring them). One idea would be to group pattern into formal groups and then allowing to specify a partial order on these groups, which is then used to specify application order of patterns. That would solve out case where NVVM pattern are meant to shadow LLVM pattern.

For now, should we land this?

Yes if it fixes the build.

mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h

LLVM::LLVMFuncOp appendOrGetFuncOp(StringRef funcName,
     mlir::OpBuilder b(op->getParentOfType<LLVMFuncOp>());
     return b.create<LLVMFuncOp>(op->getLoc(), funcName, funcType);
   }
 
   const std::string f32Func;
   const std::string f64Func;
 };
 
+namespace gpu {
+/// Returns a predicate to be used with addDynamicallyLegalOp. The predicate
+/// returns false for calls to the provided intrinsics and true otherwise.
+inline std::function<bool(Operation *)>
+filterIllegalLLVMIntrinsics(ArrayRef<StringRef> intrinsics, MLIRContext *ctx) {
+  llvm::SmallVector<StringRef, 4> illegalIds(intrinsics.begin(),
+                                             intrinsics.end());
+  return [illegalIds](Operation *op) -> bool {
+    LLVM::CallOp callOp = dyn_cast<LLVM::CallOp>(op);
+    if (!callOp || !callOp.callee())
+      return true;
+    StringRef callee = callOp.callee().getValue();
+    return !llvm::any_of(illegalIds, std::bind(StringRef::equals, callee));
+  };
+}
+} // namespace gpu
+
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_GPUCOMMON_OPTOFUNCCALLLOWERING_H_

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

void runOnOperation() override {
     populateGpuToNVVMConversionPatterns(converter, patterns);
     ConversionTarget target(getContext());
     target.addIllegalDialect<gpu::GPUDialect>();
     target.addIllegalOp<LLVM::FAbsOp, LLVM::FCeilOp, LLVM::CosOp,
                         LLVM::ExpOp>();
     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)))
       signalPassFailure();
   }
 };
 
 } // anonymous namespace

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

patterns.insert<OpToFuncCallLowering<ExpOp>>(converter, "__ocml_exp_f32",
                                                  "__ocml_exp_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.addDynamicallyLegalOp<FuncOp>(
-        [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); });
+    target.addDynamicallyLegalOp<LLVM::CallOp>(
+        gpu::filterIllegalLLVMIntrinsics({"tanh", "tanhf"}, m.getContext()));
+    target.addIllegalOp<FuncOp>();
     if (failed(applyPartialConversion(m, target, patterns, &converter)))
       signalPassFailure();
   }
 };
 
 } // anonymous namespace
 
 std::unique_ptr<OpPassBase<gpu::GPUModuleOp>>
 mlir::createLowerGpuOpsToROCDLOpsPass() {
   return std::make_unique<LowerGpuOpsToROCDLOpsPass>();
 }
 
 static PassRegistration<LowerGpuOpsToROCDLOpsPass>
     pass("convert-gpu-to-rocdl",
          "Generate ROCDL operations for gpu operations");