[mlir][spirv] Convert memref.alloca to spirv.Variable

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

 //===- MemRefToSPIRV.cpp - MemRef to SPIR-V Patterns ----------------------===//
 //
 // 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 patterns to convert MemRef dialect to SPIR-V dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "memref-to-spirv-pattern"
 
 using namespace mlir;
 

 static Value shiftValue(Location loc, Value value, Value offset, Value mask,
                         int targetBits, OpBuilder &builder) {
   Type targetType = builder.getIntegerType(targetBits);
   Value result = builder.create<spirv::BitwiseAndOp>(loc, value, mask);
   return builder.create<spirv::ShiftLeftLogicalOp>(loc, targetType, result,
                                                    offset);
 }
 
-/// Returns true if the allocations of type `t` can be lowered to SPIR-V.
-static bool isAllocationSupported(MemRefType t) {
-  // Currently only support workgroup local memory allocations with static
-  // shape and int or float or vector of int or float element type.
-  if (!(t.hasStaticShape() &&
-        SPIRVTypeConverter::getMemorySpaceForStorageClass(
-            spirv::StorageClass::Workgroup) == t.getMemorySpaceAsInt()))
+/// Returns true if the allocations of memref `type` generated from `allocOp`
+/// can be lowered to SPIR-V.
+static bool isAllocationSupported(Operation *allocOp, MemRefType type) {
+  if (isa<memref::AllocOp, memref::DeallocOp>(allocOp)) {
+    if (SPIRVTypeConverter::getMemorySpaceForStorageClass(
+            spirv::StorageClass::Workgroup) != type.getMemorySpaceAsInt())
       return false;
-  Type elementType = t.getElementType();
+  } else if (isa<memref::AllocaOp>(allocOp)) {
+    if (SPIRVTypeConverter::getMemorySpaceForStorageClass(
+            spirv::StorageClass::Function) != type.getMemorySpaceAsInt())
+      return false;
+  } else {
+    return false;
+  }
+
+  // Currently only support static shape and int or float or vector of int or
+  // float element type.
+  if (!type.hasStaticShape())
+    return false;
+
+  Type elementType = type.getElementType();
   if (auto vecType = elementType.dyn_cast<VectorType>())
     elementType = vecType.getElementType();
   return elementType.isIntOrFloat();
 }
 
 /// Returns the scope to use for atomic operations use for emulating store
 /// operations of unsupported integer bitwidths, based on the memref
 /// type. Returns None on failure.
-static Optional<spirv::Scope> getAtomicOpScope(MemRefType t) {
+static Optional<spirv::Scope> getAtomicOpScope(MemRefType type) {
   Optional<spirv::StorageClass> storageClass =
       SPIRVTypeConverter::getStorageClassForMemorySpace(
-          t.getMemorySpaceAsInt());
+          type.getMemorySpaceAsInt());
   if (!storageClass)
     return {};
   switch (*storageClass) {
   case spirv::StorageClass::StorageBuffer:
     return spirv::Scope::Device;
   case spirv::StorageClass::Workgroup:
     return spirv::Scope::Workgroup;
   default: {

 //===----------------------------------------------------------------------===//
 
 // Note that DRR cannot be used for the patterns in this file: we may need to
 // convert type along the way, which requires ConversionPattern. DRR generates
 // normal RewritePattern.
 
 namespace {
 
+/// Converts memref.alloca to SPIR-V Function variables.
+class AllocaOpPattern final : public OpConversionPattern<memref::AllocaOp> {
+public:
+  using OpConversionPattern<memref::AllocaOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::AllocaOp allocaOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 /// Converts an allocation operation to SPIR-V. Currently only supports lowering
 /// to Workgroup memory when the size is constant.  Note that this pattern needs
 /// to be applied in a pass that runs at least at spv.module scope since it wil
 /// ladd global variables into the spv.module.
 class AllocOpPattern final : public OpConversionPattern<memref::AllocOp> {
 public:
   using OpConversionPattern<memref::AllocOp>::OpConversionPattern;
 
public:
   LogicalResult
   matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
+// AllocaOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+AllocaOpPattern::matchAndRewrite(memref::AllocaOp allocaOp, OpAdaptor adaptor,
+                                 ConversionPatternRewriter &rewriter) const {
+  MemRefType allocType = allocaOp.getType();
+  if (!isAllocationSupported(allocaOp, allocType))
+    return rewriter.notifyMatchFailure(allocaOp, "unhandled allocation type");
+
+  // Get the SPIR-V type for the allocation.
+  Type spirvType = getTypeConverter()->convertType(allocType);
+  rewriter.replaceOpWithNewOp<spirv::VariableOp>(allocaOp, spirvType,
+                                                 spirv::StorageClass::Function,
+                                                 /*initializer=*/nullptr);
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // AllocOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 AllocOpPattern::matchAndRewrite(memref::AllocOp operation, OpAdaptor adaptor,
                                 ConversionPatternRewriter &rewriter) const {
   MemRefType allocType = operation.getType();
-  if (!isAllocationSupported(allocType))
-    return operation.emitError("unhandled allocation type");
+  if (!isAllocationSupported(operation, allocType))
+    return rewriter.notifyMatchFailure(operation, "unhandled allocation type");
 
   // Get the SPIR-V type for the allocation.
   Type spirvType = getTypeConverter()->convertType(allocType);
 
   // Insert spv.GlobalVariable for this allocation.
   Operation *parent =
       SymbolTable::getNearestSymbolTable(operation->getParentOp());
   if (!parent)

 // DeallocOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 DeallocOpPattern::matchAndRewrite(memref::DeallocOp operation,
                                   OpAdaptor adaptor,
                                   ConversionPatternRewriter &rewriter) const {
   MemRefType deallocType = operation.memref().getType().cast<MemRefType>();
-  if (!isAllocationSupported(deallocType))
-    return operation.emitError("unhandled deallocation type");
+  if (!isAllocationSupported(operation, deallocType))
+    return rewriter.notifyMatchFailure(operation, "unhandled allocation type");
   rewriter.eraseOp(operation);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // LoadOp
 //===----------------------------------------------------------------------===//
 

 
 //===----------------------------------------------------------------------===//
 // Pattern population
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 void populateMemRefToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                    RewritePatternSet &patterns) {
-  patterns.add<AllocOpPattern, DeallocOpPattern, IntLoadOpPattern,
-               IntStoreOpPattern, LoadOpPattern, StoreOpPattern>(
+  patterns
+      .add<AllocaOpPattern, AllocOpPattern, DeallocOpPattern, IntLoadOpPattern,
           typeConverter, patterns.getContext());
 }
 } // namespace mlir

mlir/test/Conversion/MemRefToSPIRV/alloc.mlir

 
 // -----
 
 module attributes {
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, {}>
   }
 {
-  func.func @alloc_dealloc_dynamic_workgroup_mem(%arg0 : index) {
-    // expected-error @+1 {{unhandled allocation type}}
+  // CHECK-LABEL: func @alloc_dynamic_size
+  func.func @alloc_dynamic_size(%arg0 : index) -> f32 {
+    // CHECK: memref.alloc
     %0 = memref.alloc(%arg0) : memref<4x?xf32, 3>
-    return
+    %1 = memref.load %0[%arg0, %arg0] : memref<4x?xf32, 3>
+    return %1: f32
   }
 }
 
 // -----
 
 module attributes {
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, {}>
   }
 {
-  func.func @alloc_dealloc_mem() {
-    // expected-error @+1 {{unhandled allocation type}}
+  // CHECK-LABEL: func @alloc_unsupported_memory_space
+  func.func @alloc_unsupported_memory_space(%arg0: index) -> f32 {
+    // CHECK: memref.alloc
     %0 = memref.alloc() : memref<4x5xf32>
-    return
+    %1 = memref.load %0[%arg0, %arg0] : memref<4x5xf32>
+    return %1: f32
   }
 }
 
 
 // -----
 
 module attributes {
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, {}>
   }
 {
-  func.func @alloc_dealloc_dynamic_workgroup_mem(%arg0 : memref<4x?xf32, 3>) {
-    // expected-error @+1 {{unhandled deallocation type}}
+  // CHECK-LABEL: func @dealloc_dynamic_size
+  func.func @dealloc_dynamic_size(%arg0 : memref<4x?xf32, 3>) {
+    // CHECK: memref.dealloc
     memref.dealloc %arg0 : memref<4x?xf32, 3>
     return
   }
 }
 
 // -----
 
 module attributes {
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, {}>
   }
 {
-  func.func @alloc_dealloc_mem(%arg0 : memref<4x5xf32>) {
-    // expected-error @+1 {{unhandled deallocation type}}
+  // CHECK-LABEL: func @dealloc_unsupported_memory_space
+  func.func @dealloc_unsupported_memory_space(%arg0 : memref<4x5xf32>) {
+    // CHECK: memref.dealloc
     memref.dealloc %arg0 : memref<4x5xf32>
     return
   }
 }

mlir/test/Conversion/MemRefToSPIRV/alloca.mlir

+// RUN: mlir-opt -split-input-file -convert-memref-to-spirv -canonicalize -verify-diagnostics %s -o - | FileCheck %s
+
+module attributes {spv.target_env = #spv.target_env<#spv.vce<v1.3, [Shader], []>, {}>} {
+  func.func @alloc_function_variable(%arg0 : index, %arg1 : index) {
+    %0 = memref.alloca() : memref<4x5xf32, 6>
+    %1 = memref.load %0[%arg0, %arg1] : memref<4x5xf32, 6>
+    memref.store %1, %0[%arg0, %arg1] : memref<4x5xf32, 6>
+    return
+  }
+}
+
+// CHECK-LABEL: func @alloc_function_variable
+//       CHECK:   %[[VAR:.+]] = spv.Variable : !spv.ptr<!spv.struct<(!spv.array<20 x f32, stride=4>)>, Function>
+//       CHECK:   %[[LOADPTR:.+]] = spv.AccessChain %[[VAR]]
+//       CHECK:   %[[VAL:.+]] = spv.Load "Function" %[[LOADPTR]] : f32
+//       CHECK:   %[[STOREPTR:.+]] = spv.AccessChain %[[VAR]]
+//       CHECK:   spv.Store "Function" %[[STOREPTR]], %[[VAL]] : f32
+
+
+// -----
+
+module attributes {spv.target_env = #spv.target_env<#spv.vce<v1.3, [Shader], []>, {}>} {
+  func.func @two_allocs() {
+    %0 = memref.alloca() : memref<4x5xf32, 6>
+    %1 = memref.alloca() : memref<2x3xi32, 6>
+    return
+  }
+}
+
+// CHECK-LABEL: func @two_allocs
+//   CHECK-DAG: spv.Variable : !spv.ptr<!spv.struct<(!spv.array<6 x i32, stride=4>)>, Function>
+//   CHECK-DAG: spv.Variable : !spv.ptr<!spv.struct<(!spv.array<20 x f32, stride=4>)>, Function>
+
+// -----
+
+module attributes {spv.target_env = #spv.target_env<#spv.vce<v1.3, [Shader], []>, {}>} {
+  func.func @two_allocs_vector() {
+    %0 = memref.alloca() : memref<4xvector<4xf32>, 6>
+    %1 = memref.alloca() : memref<2xvector<2xi32>, 6>
+    return
+  }
+}
+
+// CHECK-LABEL: func @two_allocs_vector
+//   CHECK-DAG: spv.Variable : !spv.ptr<!spv.struct<(!spv.array<2 x vector<2xi32>, stride=8>)>, Function>
+//   CHECK-DAG: spv.Variable : !spv.ptr<!spv.struct<(!spv.array<4 x vector<4xf32>, stride=16>)>, Function>
+
+
+// -----
+
+module attributes {spv.target_env = #spv.target_env<#spv.vce<v1.3, [Shader], []>, {}>} {
+  // CHECK-LABEL: func @alloc_dynamic_size
+  func.func @alloc_dynamic_size(%arg0 : index) -> f32 {
+    // CHECK: memref.alloca
+    %0 = memref.alloca(%arg0) : memref<4x?xf32, 6>
+    %1 = memref.load %0[%arg0, %arg0] : memref<4x?xf32, 6>
+    return %1: f32
+  }
+}
+
+// -----
+
+module attributes {spv.target_env = #spv.target_env<#spv.vce<v1.3, [Shader], []>, {}>} {
+  // CHECK-LABEL: func @alloc_unsupported_memory_space
+  func.func @alloc_unsupported_memory_space(%arg0: index) -> f32 {
+    // CHECK: memref.alloca
+    %0 = memref.alloca() : memref<4x5xf32>
+    %1 = memref.load %0[%arg0, %arg0] : memref<4x5xf32>
+    return %1: f32
+  }
+}