[mlir][spirv] Fix nullptr dereference in UnifyAliasedResource

[kuhar]

Summary

Fixes: https://github.com/llvm/llvm-project/issues/62368

mlir/lib/Dialect/SPIRV/Transforms/UnifyAliasedResourcePass.cpp

ResourceAliasAnalysis::ResourceAliasAnalysis(Operation *root) {
   // For each resource set, analyze whether we can unify; if so, try to identify
   // a canonical resource, whose element type has the largest bitwidth.
   for (const auto &descriptorResource : aliasedResources) {
     recordIfUnifiable(descriptorResource.first, descriptorResource.second);
   }
 }
 
 bool ResourceAliasAnalysis::shouldUnify(Operation *op) const {
+  if (!op)
+    return false;
+
   if (auto varOp = dyn_cast<spirv::GlobalVariableOp>(op)) {
     auto canonicalOp = getCanonicalResource(varOp);
     return canonicalOp && varOp != canonicalOp;
   }
   if (auto addressOp = dyn_cast<spirv::AddressOfOp>(op)) {
     auto moduleOp = addressOp->getParentOfType<spirv::ModuleOp>();
     auto *varOp =
         SymbolTable::lookupSymbolIn(moduleOp, addressOp.getVariable());
public:
   explicit UnifyAliasedResourcePass(spirv::GetTargetEnvFn getTargetEnv)
       : getTargetEnvFn(std::move(getTargetEnv)) {}
 
   void runOnOperation() override;
 
 private:
   spirv::GetTargetEnvFn getTargetEnvFn;
 };
-} // namespace
 
 void UnifyAliasedResourcePass::runOnOperation() {
   spirv::ModuleOp moduleOp = getOperation();
   MLIRContext *context = &getContext();
 
   if (getTargetEnvFn) {
-    // This pass is only needed for targeting WebGPU, Metal, or layering Vulkan
-    // on Metal via MoltenVK, where we need to translate SPIR-V into WGSL or
-    // MSL. The translation has limitations.
+    // This pass is only needed for targeting WebGPU, Metal, or layering
+    // Vulkan on Metal via MoltenVK, where we need to translate SPIR-V into
+    // WGSL or MSL. The translation has limitations.
     spirv::TargetEnvAttr targetEnv = getTargetEnvFn(moduleOp);
     spirv::ClientAPI clientAPI = targetEnv.getClientAPI();
     bool isVulkanOnAppleDevices =
         clientAPI == spirv::ClientAPI::Vulkan &&
         targetEnv.getVendorID() == spirv::Vendor::Apple;
     if (clientAPI != spirv::ClientAPI::WebGPU &&
         clientAPI != spirv::ClientAPI::Metal && !isVulkanOnAppleDevices)
       return;
void UnifyAliasedResourcePass::runOnOperation() {
   AliasedResourceMap resourceMap =
       collectAliasedResources(cast<spirv::ModuleOp>(moduleOp));
   for (const auto &dr : resourceMap) {
     const auto &resources = dr.second;
     if (resources.size() == 1)
       resources.front()->removeAttr("aliased");
   }
 }
+} // namespace
 
 std::unique_ptr<mlir::OperationPass<spirv::ModuleOp>>
 spirv::createUnifyAliasedResourcePass(spirv::GetTargetEnvFn getTargetEnv) {
   return std::make_unique<UnifyAliasedResourcePass>(std::move(getTargetEnv));
 }

mlir/test/Dialect/SPIRV/Transforms/unify-aliased-resource.mlir

 
 // CHECK:   %[[LD0:.+]] = spirv.Load "StorageBuffer" %{{.+}} : vector<4xf16>
 // CHECK:   %[[LD1:.+]] = spirv.Load "StorageBuffer" %{{.+}} : vector<4xf16>
 // CHECK:   %[[BC0:.+]] = spirv.Bitcast %[[LD0]] : vector<4xf16> to vector<2xf32>
 // CHECK:   %[[BC1:.+]] = spirv.Bitcast %[[LD1]] : vector<4xf16> to vector<2xf32>
 // CHECK:   %[[CC:.+]] = spirv.CompositeConstruct %[[BC0]], %[[BC1]] : (vector<2xf32>, vector<2xf32>) -> vector<4xf32>
 // CHECK:   spirv.ReturnValue %[[CC]]
 
+// -----
+
+// Make sure we do not crash on function arguments.
+
+spirv.module Logical GLSL450 {
+  spirv.func @main(%arg0: !spirv.ptr<!spirv.struct<(!spirv.rtarray<f32, stride=4> [0])>, StorageBuffer>) "None" {
+    %cst0_i32 = spirv.Constant 0 : i32
+    %0 = spirv.AccessChain %arg0[%cst0_i32, %cst0_i32] : !spirv.ptr<!spirv.struct<(!spirv.rtarray<f32, stride=4> [0])>, StorageBuffer>, i32, i32
+    spirv.Return
+  }
+}
+
+// CHECK-LABEL: spirv.module
+// CHECK-LABEL: spirv.func @main
+// CHECK-SAME:  (%{{.+}}: !spirv.ptr<!spirv.struct<(!spirv.rtarray<f32, stride=4> [0])>, StorageBuffer>) "None"
+// CHECK:       spirv.Return