83 lines · cpp
1//===- Utils.cpp - Transform utilities ------------------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Dialect/NVGPU/Transforms/Utils.h"10 11#include "mlir/Dialect/MemRef/IR/MemRef.h"12#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"13#include "mlir/Dialect/Vector/IR/VectorOps.h"14 15using namespace mlir;16using namespace mlir::nvgpu;17 18Operation::operand_range nvgpu::getIndices(Operation *op) {19 if (auto ldmatrixOp = dyn_cast<LdMatrixOp>(op))20 return ldmatrixOp.getIndices();21 if (auto copyOp = dyn_cast<DeviceAsyncCopyOp>(op))22 return copyOp.getDstIndices();23 if (auto loadOp = dyn_cast<memref::LoadOp>(op))24 return loadOp.getIndices();25 if (auto storeOp = dyn_cast<memref::StoreOp>(op))26 return storeOp.getIndices();27 if (auto vectorReadOp = dyn_cast<vector::LoadOp>(op))28 return vectorReadOp.getIndices();29 if (auto vectorStoreOp = dyn_cast<vector::StoreOp>(op))30 return vectorStoreOp.getIndices();31 if (auto transferReadOp = dyn_cast<vector::TransferReadOp>(op))32 return transferReadOp.getIndices();33 if (auto transferWriteOp = dyn_cast<vector::TransferWriteOp>(op))34 return transferWriteOp.getIndices();35 llvm_unreachable("unsupported op type");36}37 38void nvgpu::setIndices(Operation *op, ArrayRef<Value> indices) {39 if (auto ldmatrixOp = dyn_cast<LdMatrixOp>(op))40 return ldmatrixOp.getIndicesMutable().assign(indices);41 if (auto copyOp = dyn_cast<DeviceAsyncCopyOp>(op))42 return copyOp.getDstIndicesMutable().assign(indices);43 if (auto loadOp = dyn_cast<memref::LoadOp>(op))44 return loadOp.getIndicesMutable().assign(indices);45 if (auto storeOp = dyn_cast<memref::StoreOp>(op))46 return storeOp.getIndicesMutable().assign(indices);47 if (auto vectorReadOp = dyn_cast<vector::LoadOp>(op))48 return vectorReadOp.getIndicesMutable().assign(indices);49 if (auto vectorStoreOp = dyn_cast<vector::StoreOp>(op))50 return vectorStoreOp.getIndicesMutable().assign(indices);51 if (auto transferReadOp = dyn_cast<vector::TransferReadOp>(op))52 return transferReadOp.getIndicesMutable().assign(indices);53 if (auto transferWriteOp = dyn_cast<vector::TransferWriteOp>(op))54 return transferWriteOp.getIndicesMutable().assign(indices);55 llvm_unreachable("unsupported op type");56}57 58Value nvgpu::getValueStored(Operation *op) {59 if (auto storeOp = dyn_cast<memref::StoreOp>(op))60 return storeOp.getValueToStore();61 if (auto transferWrite = dyn_cast<vector::TransferWriteOp>(op))62 return transferWrite.getValue();63 if (auto storeOp = dyn_cast<vector::StoreOp>(op))64 return storeOp.getValueToStore();65 llvm_unreachable("unsupported op type");66}67 68Value nvgpu::getMemrefOperand(Operation *op) {69 if (auto loadOp = dyn_cast<memref::LoadOp>(op))70 return loadOp.getMemref();71 if (auto storeOp = dyn_cast<memref::StoreOp>(op))72 return storeOp.getMemref();73 if (auto transferWrite = dyn_cast<vector::TransferWriteOp>(op))74 return transferWrite.getBase();75 if (auto transferRead = dyn_cast<vector::TransferReadOp>(op))76 return transferRead.getBase();77 if (auto storeOp = dyn_cast<vector::StoreOp>(op))78 return storeOp.getBase();79 if (auto loadOp = dyn_cast<vector::LoadOp>(op))80 return loadOp.getBase();81 llvm_unreachable("unsupported op type");82}83