Add interface for get and set indices functions

Author: erman-gurses

mlir/lib/Dialect/AMDGPU/Transforms/Utils.cpp

@@ -1,11 +1,3 @@
-//===- Utils.cpp - Transform utilities ------------------------------------===//
-//
-// 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/AMDGPU/Transforms/Utils.h"
 
 #include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
@@ -15,33 +7,66 @@
 using namespace mlir;
 using namespace mlir::amdgpu;
 
+// Define an interface for operations with indices
+class IndicesInterface {
+public:
+  virtual std::optional<Operation::operand_range> getIndices() = 0;
+  virtual void setIndices(ArrayRef<Value> indices) = 0;
+  virtual ~IndicesInterface() = default;
+};
+
+// Implement a generic class that uses IndicesInterface
+class OperationWithIndices : public IndicesInterface {
+private:
+  Operation *op;
+  template <typename OpType>
+  static std::optional<Operation::operand_range> getIndicesImpl(Operation *op) {
+    if (auto specificOp = dyn_cast<OpType>(op))
+      return specificOp.getIndices();
+    return std::nullopt;
+  }
+
+  template <typename OpType>
+  static void setIndicesImpl(Operation *op, ArrayRef<Value> indices) {
+    if (auto specificOp = dyn_cast<OpType>(op))
+      specificOp.getIndicesMutable().assign(indices);
+  }
+
+public:
+  OperationWithIndices(Operation *op) : op(op) {}
+
+  std::optional<Operation::operand_range> getIndices() override {
+    auto result = getIndicesImpl<memref::LoadOp>(op);
+    if (!result)
+      result = getIndicesImpl<memref::StoreOp>(op);
+    if (!result)
+      result = getIndicesImpl<vector::LoadOp>(op);
+    if (!result)
+      result = getIndicesImpl<vector::StoreOp>(op);
+    if (!result)
+      result = getIndicesImpl<vector::TransferReadOp>(op);
+    if (!result)
+      result = getIndicesImpl<vector::TransferWriteOp>(op);
+
+    return result;
+  }
+
+  void setIndices(ArrayRef<Value> indices) override {
+    setIndicesImpl<memref::LoadOp>(op, indices);
+    setIndicesImpl<memref::StoreOp>(op, indices);
+    setIndicesImpl<vector::LoadOp>(op, indices);
+    setIndicesImpl<vector::StoreOp>(op, indices);
+    setIndicesImpl<vector::TransferReadOp>(op, indices);
+    setIndicesImpl<vector::TransferWriteOp>(op, indices);
+  }
+};
+
 std::optional<Operation::operand_range> amdgpu::getIndices(Operation *op) {
-  if (auto loadOp = dyn_cast<memref::LoadOp>(op))
-    return loadOp.getIndices();
-  if (auto storeOp = dyn_cast<memref::StoreOp>(op))
-    return storeOp.getIndices();
-  if (auto vectorReadOp = dyn_cast<vector::LoadOp>(op))
-    return vectorReadOp.getIndices();
-  if (auto vectorStoreOp = dyn_cast<vector::StoreOp>(op))
-    return vectorStoreOp.getIndices();
-  if (auto transferReadOp = dyn_cast<vector::TransferReadOp>(op))
-    return transferReadOp.getIndices();
-  if (auto transferWriteOp = dyn_cast<vector::TransferWriteOp>(op))
-    return transferWriteOp.getIndices();
-  return std::nullopt;
+  OperationWithIndices operationWithIndices(op);
+  return operationWithIndices.getIndices();
 }
 
 void amdgpu::setIndices(Operation *op, ArrayRef<Value> indices) {
-  if (auto loadOp = dyn_cast<memref::LoadOp>(op))
-    return loadOp.getIndicesMutable().assign(indices);
-  if (auto storeOp = dyn_cast<memref::StoreOp>(op))
-    return storeOp.getIndicesMutable().assign(indices);
-  if (auto vectorReadOp = dyn_cast<vector::LoadOp>(op))
-    return vectorReadOp.getIndicesMutable().assign(indices);
-  if (auto vectorStoreOp = dyn_cast<vector::StoreOp>(op))
-    return vectorStoreOp.getIndicesMutable().assign(indices);
-  if (auto transferReadOp = dyn_cast<vector::TransferReadOp>(op))
-    return transferReadOp.getIndicesMutable().assign(indices);
-  if (auto transferWriteOp = dyn_cast<vector::TransferWriteOp>(op))
-    return transferWriteOp.getIndicesMutable().assign(indices);
-}
+  OperationWithIndices operationWithIndices(op);
+  operationWithIndices.setIndices(indices);
+}