[spirv] NFC: Replace OwningSPIRVModuleRef with OwningOpRef

[rriddle]

Thanks Lei!

mlir/include/mlir/Dialect/SPIRV/Linking/ModuleCombiner.h

68–69

Can you just add an include for LLVM.h?

mlir/lib/Target/SPIRV/Deserialization/Deserialization.cpp

12–23

[antiagainst]

Address comments

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVModule.h

-//===- SPIRVModule.h - SPIR-V Module Utilities ------------------*- C++ -*-===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MLIR_DIALECT_SPIRV_IR_SPIRVMODULE_H
-#define MLIR_DIALECT_SPIRV_IR_SPIRVMODULE_H
-
-#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
-#include "mlir/IR/OwningOpRef.h"
-
-namespace mlir {
-namespace spirv {
-
-/// This class acts as an owning reference to a SPIR-V module, and will
-/// automatically destroy the held module on destruction if the held module
-/// is valid.
-// TODO: Remove this class in favor of using OwningOpRef directly.
-class OwningSPIRVModuleRef : public OwningOpRef<spirv::ModuleOp> {
-public:
-  using OwningOpRef<spirv::ModuleOp>::OwningOpRef;
-};
-
-} // end namespace spirv
-} // end namespace mlir
-
-#endif // MLIR_DIALECT_SPIRV_IR_SPIRVMODULE_H

mlir/include/mlir/Dialect/SPIRV/Linking/ModuleCombiner.h

 //===- ModuleCombiner.h - MLIR SPIR-V Module Combiner -----------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file declares the entry point to the SPIR-V module combiner library.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPIRV_LINKING_MODULECOMBINER_H_
 #define MLIR_DIALECT_SPIRV_LINKING_MODULECOMBINER_H_
 
-#include "mlir/Dialect/SPIRV/IR/SPIRVModule.h"
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/SmallVector.h"
+#include "mlir/IR/OwningOpRef.h"
+#include "mlir/Support/LLVM.h"
 
 namespace mlir {
 class OpBuilder;
 
 namespace spirv {
 class ModuleOp;
 
 /// To combine a number of MLIR SPIR-V modules, we move all the module-level ops

 ///                           spirv::ModuleOp that contains the renamed symbol,
 ///                           a StringRef to the old symbol name, and a
 ///                           StringRef to the new symbol name. Note that it is
 ///                           the responsibility of the caller to properly
 ///                           retain the storage underlying the passed
 ///                           StringRefs if the listener callback outlives this
 ///                           function call.
 ///
 /// \return the combined module.
-OwningSPIRVModuleRef
-combine(llvm::MutableArrayRef<ModuleOp> modules,
-        OpBuilder &combinedModuleBuilder,
-        llvm::function_ref<void(ModuleOp, StringRef, StringRef)>
-            symbRenameListener);
+OwningOpRef<spirv::ModuleOp>
+combine(MutableArrayRef<ModuleOp> modules, OpBuilder &combinedModuleBuilder,
+        function_ref<void(ModuleOp, StringRef, StringRef)> symbRenameListener);
 } // namespace spirv
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPIRV_LINKING_MODULECOMBINER_H_

mlir/include/mlir/Target/SPIRV/Deserialization.h

 //===- Serialization.h - MLIR SPIR-V (De)serialization ----------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file declares the entry points for deserializing SPIR-V binary modules.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TARGET_SPIRV_DESERIALIZATION_H
 #define MLIR_TARGET_SPIRV_DESERIALIZATION_H
 
+#include "mlir/IR/OwningOpRef.h"
 #include "mlir/Support/LLVM.h"
 
 namespace mlir {
-struct LogicalResult;
 class MLIRContext;
 
 namespace spirv {
-class OwningSPIRVModuleRef;
+class ModuleOp;
 
 /// Deserializes the given SPIR-V `binary` module and creates a MLIR ModuleOp
 /// in the given `context`. Returns the ModuleOp on success; otherwise, reports
 /// errors to the error handler registered with `context` and returns a null
 /// module.
-OwningSPIRVModuleRef deserialize(ArrayRef<uint32_t> binary,
+OwningOpRef<spirv::ModuleOp> deserialize(ArrayRef<uint32_t> binary,
                                          MLIRContext *context);
 
 } // end namespace spirv
 } // end namespace mlir
 
 #endif // MLIR_TARGET_SPIRV_DESERIALIZATION_H

mlir/lib/Dialect/SPIRV/Linking/ModuleCombiner/ModuleCombiner.cpp

llvm::hash_code computeHash(Block *block) {
   return hash;
 }
 
 namespace mlir {
 namespace spirv {
 
 // TODO Properly test symbol rename listener mechanism.
 
-OwningSPIRVModuleRef
+OwningOpRef<spirv::ModuleOp>
 combine(llvm::MutableArrayRef<spirv::ModuleOp> modules,
         OpBuilder &combinedModuleBuilder,
         llvm::function_ref<void(ModuleOp, StringRef, StringRef)>
             symRenameListener) {
   unsigned lastUsedID = 0;
 
   if (modules.empty())
     return nullptr;

mlir/lib/Target/SPIRV/Deserialization/Deserialization.cpp

 //===- Deserialization.cpp - MLIR SPIR-V Deserialization ------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Target/SPIRV/Deserialization.h"
 
 #include "Deserializer.h"
 
-namespace mlir {
-spirv::OwningSPIRVModuleRef spirv::deserialize(ArrayRef<uint32_t> binary,
+using namespace mlir;
+
+OwningOpRef<spirv::ModuleOp> spirv::deserialize(ArrayRef<uint32_t> binary,
                                                 MLIRContext *context) {
   Deserializer deserializer(binary, context);
 
   if (failed(deserializer.deserialize()))
     return nullptr;
 
   return deserializer.collect();
 }
-} // namespace mlir

mlir/lib/Target/SPIRV/Deserialization/Deserializer.h

 //===- Deserializer.h - MLIR SPIR-V Deserializer ----------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file declares the SPIR-V binary to MLIR SPIR-V module deserializer.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TARGET_SPIRV_DESERIALIZER_H
 #define MLIR_TARGET_SPIRV_DESERIALIZER_H
 
 #include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
-#include "mlir/Dialect/SPIRV/IR/SPIRVModule.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/IR/Builders.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringRef.h"
 #include <cstdint>
 
 //===----------------------------------------------------------------------===//
public:
   /// Creates a deserializer for the given SPIR-V `binary` module.
   /// The SPIR-V ModuleOp will be created into `context.
   explicit Deserializer(ArrayRef<uint32_t> binary, MLIRContext *context);
 
   /// Deserializes the remembered SPIR-V binary module.
   LogicalResult deserialize();
 
   /// Collects the final SPIR-V ModuleOp.
-  spirv::OwningSPIRVModuleRef collect();
+  OwningOpRef<spirv::ModuleOp> collect();
 
 private:
   //===--------------------------------------------------------------------===//
   // Module structure
   //===--------------------------------------------------------------------===//
 
   /// Initializes the `module` ModuleOp in this deserializer instance.
-  spirv::OwningSPIRVModuleRef createModuleOp();
+  OwningOpRef<spirv::ModuleOp> createModuleOp();
 
   /// Processes SPIR-V module header in `binary`.
   LogicalResult processHeader();
 
   /// Processes the SPIR-V OpCapability with `operands` and updates bookkeeping
   /// in the deserializer.
   LogicalResult processCapability(ArrayRef<uint32_t> operands);
 
private:
 
   /// MLIRContext to create SPIR-V ModuleOp into.
   MLIRContext *context;
 
   // TODO: create Location subclass for binary blob
   Location unknownLoc;
 
   /// The SPIR-V ModuleOp.
-  spirv::OwningSPIRVModuleRef module;
+  OwningOpRef<spirv::ModuleOp> module;
 
   /// The current function under construction.
   Optional<spirv::FuncOp> curFunction;
 
   /// The current block under construction.
   Block *curBlock = nullptr;
 
   OpBuilder opBuilder;

mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp

 //===- Deserializer.cpp - MLIR SPIR-V Deserializer ------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the SPIR-V binary to MLIR SPIR-V module deserializer.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Deserializer.h"
 
 #include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
-#include "mlir/Dialect/SPIRV/IR/SPIRVModule.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Location.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Target/SPIRV/SPIRVBinaryUtils.h"
 #include "llvm/ADT/STLExtras.h"
LogicalResult spirv::Deserializer::deserialize() {
   }
 
   attachVCETriple();
 
   LLVM_DEBUG(llvm::dbgs() << "+++ completed deserialization +++\n");
   return success();
 }
 
-spirv::OwningSPIRVModuleRef spirv::Deserializer::collect() {
+OwningOpRef<spirv::ModuleOp> spirv::Deserializer::collect() {
   return std::move(module);
 }
 
 //===----------------------------------------------------------------------===//
 // Module structure
 //===----------------------------------------------------------------------===//
 
-spirv::OwningSPIRVModuleRef spirv::Deserializer::createModuleOp() {
+OwningOpRef<spirv::ModuleOp> spirv::Deserializer::createModuleOp() {
   OpBuilder builder(context);
   OperationState state(unknownLoc, spirv::ModuleOp::getOperationName());
   spirv::ModuleOp::build(builder, state);
   return cast<spirv::ModuleOp>(Operation::create(state));
 }
 
 LogicalResult spirv::Deserializer::processHeader() {
   if (binary.size() < spirv::kHeaderWordCount)

mlir/lib/Target/SPIRV/TranslateRegistration.cpp

 //===- TranslateRegistration.cpp - hooks to mlir-translate ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a translation from SPIR-V binary module to MLIR SPIR-V
 // ModuleOp.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
-#include "mlir/Dialect/SPIRV/IR/SPIRVModule.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/Parser.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Target/SPIRV/Deserialization.h"
 #include "mlir/Target/SPIRV/Serialization.h"
if (size % sizeof(uint32_t) != 0) {
     emitError(UnknownLoc::get(context))
         << "SPIR-V binary module must contain integral number of 32-bit words";
     return {};
   }
 
   auto binary = llvm::makeArrayRef(reinterpret_cast<const uint32_t *>(start),
                                    size / sizeof(uint32_t));
 
-  spirv::OwningSPIRVModuleRef spirvModule = spirv::deserialize(binary, context);
+  OwningOpRef<spirv::ModuleOp> spirvModule =
+      spirv::deserialize(binary, context);
   if (!spirvModule)
     return {};
 
   OwningModuleRef module(ModuleOp::create(FileLineColLoc::get(
       context, input->getBufferIdentifier(), /*line=*/0, /*column=*/0)));
   module->getBody()->push_front(spirvModule.release());
 
   return module;
static LogicalResult roundTripModule(ModuleOp srcModule, bool emitDebugInfo,
 
   if (failed(spirv::serialize(*spirvModules.begin(), binary, emitDebugInfo)))
     return failure();
 
   MLIRContext deserializationContext(context->getDialectRegistry());
   // TODO: we should only load the required dialects instead of all dialects.
   deserializationContext.loadAllAvailableDialects();
   // Then deserialize to get back a SPIR-V module.
-  spirv::OwningSPIRVModuleRef spirvModule =
+  OwningOpRef<spirv::ModuleOp> spirvModule =
       spirv::deserialize(binary, &deserializationContext);
   if (!spirvModule)
     return failure();
 
   // Wrap around in a new MLIR module.
   OwningModuleRef dstModule(ModuleOp::create(
       FileLineColLoc::get(&deserializationContext,
                           /*filename=*/"", /*line=*/0, /*column=*/0)));

mlir/test/lib/Dialect/SPIRV/TestGLSLCanonicalization.cpp

 //===- TestGLSLCanonicalization.cpp - Pass to test GLSL-specific pattterns ===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SPIRV/IR/SPIRVGLSLCanonicalization.h"
-#include "mlir/Dialect/SPIRV/IR/SPIRVModule.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 
 namespace {
 class TestGLSLCanonicalizationPass
     : public PassWrapper<TestGLSLCanonicalizationPass,

mlir/test/lib/Dialect/SPIRV/TestModuleCombiner.cpp

class TestModuleCombinerPass
     : public PassWrapper<TestModuleCombinerPass,
                          OperationPass<mlir::ModuleOp>> {
 public:
   TestModuleCombinerPass() = default;
   TestModuleCombinerPass(const TestModuleCombinerPass &) {}
   void runOnOperation() override;
 
 private:
-  mlir::spirv::OwningSPIRVModuleRef combinedModule;
+  OwningOpRef<spirv::ModuleOp> combinedModule;
 };
 } // namespace
 
 void TestModuleCombinerPass::runOnOperation() {
   auto modules = llvm::to_vector<4>(getOperation().getOps<spirv::ModuleOp>());
 
   OpBuilder combinedModuleBuilder(modules[0]);
   combinedModule = spirv::combine(modules, combinedModuleBuilder, nullptr);

mlir/unittests/Dialect/SPIRV/DeserializationTest.cpp

 //===- DeserializationTest.cpp - SPIR-V Deserialization Tests -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // The purpose of this file is to provide negative deserialization tests.
 // For positive deserialization tests, please use serialization and
 // deserialization for roundtripping.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Target/SPIRV/Deserialization.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
-#include "mlir/Dialect/SPIRV/IR/SPIRVModule.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Target/SPIRV/SPIRVBinaryUtils.h"
 #include "gmock/gmock.h"
 
 #include <memory>
 
DeserializationTest() {
     // Register a diagnostic handler to capture the diagnostic so that we can
     // check it later.
     context.getDiagEngine().registerHandler([&](Diagnostic &diag) {
       diagnostic.reset(new Diagnostic(std::move(diag)));
     });
   }
 
   /// Performs deserialization and returns the constructed spv.module op.
-  spirv::OwningSPIRVModuleRef deserialize() {
+  OwningOpRef<spirv::ModuleOp> deserialize() {
     return spirv::deserialize(binary, &context);
   }
 
   /// Checks there is a diagnostic generated with the given `errorMessage`.
   void expectDiagnostic(StringRef errorMessage) {
     ASSERT_NE(nullptr, diagnostic.get());
 
     // TODO: check error location too.

mlir/unittests/Dialect/SPIRV/SerializationTest.cpp

 //===- SerializationTest.cpp - SPIR-V Serialization Tests -----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains corner case tests for the SPIR-V serializer that are not
 // covered by normal serialization and deserialization roundtripping.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Target/SPIRV/Serialization.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
-#include "mlir/Dialect/SPIRV/IR/SPIRVModule.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Target/SPIRV/SPIRVBinaryUtils.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
while (currOffset < binarySize) {
       }
       currOffset += wordCount;
     }
     return false;
   }
 
 protected:
   MLIRContext context;
-  spirv::OwningSPIRVModuleRef module;
+  OwningOpRef<spirv::ModuleOp> module;
   SmallVector<uint32_t, 0> binary;
 };
 
 //===----------------------------------------------------------------------===//
 // Block decoration
 //===----------------------------------------------------------------------===//
 
 TEST_F(SerializationTest, BlockDecorationTest) {