From 53ac5126ed83e6b535030f60251ce08753cafcb6 Mon Sep 17 00:00:00 2001
From: ravil-mobile <ravil.aviva.com@gmail.com>
Date: Tue, 25 Jun 2024 13:43:24 +0000
Subject: [PATCH] [AMD] quick fix of `BypassEpilogueSMEM` rewrite

`val.getDefiningOp()` can return `nullptr`. In this case,
we must fail the `BypassEpilogueSMEM` rewrite for
the given op. This prevents run-time crashes.
---
 third_party/amd/lib/TritonAMDGPUTransforms/OptimizeEpilogue.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/third_party/amd/lib/TritonAMDGPUTransforms/OptimizeEpilogue.cpp b/third_party/amd/lib/TritonAMDGPUTransforms/OptimizeEpilogue.cpp
index 6c6475b754db..7d24bf758172 100644
--- a/third_party/amd/lib/TritonAMDGPUTransforms/OptimizeEpilogue.cpp
+++ b/third_party/amd/lib/TritonAMDGPUTransforms/OptimizeEpilogue.cpp
@@ -100,6 +100,8 @@ class BypassEpilogueSMEM : public mlir::RewritePattern {
     llvm::SmallVector<mlir::Operation *> chainedOps;
     while (true) {
       auto chainedOp = val.getDefiningOp();
+      if (!chainedOp)
+        return mlir::failure();
       if (llvm::isa<triton::gpu::ConvertLayoutOp>(chainedOp))
         break;
       if (!chainedOp->hasOneUse())