[mlir][spirv] Add lowering of std alloca to SPIR-V

[ThomasRaoux]

Summary

alloca in function or private address space can be lowered to VariableOp. This allow values in memory if needed.

[hanchung]

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

306–307

I think we can rename spirvType as ptrType, and then we can remove the comment. What do you think?

311

/*initializer=*/=nullptr

[bondhugula]

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

296–300

You may have to check if the parent op has the AutomaticAllocationScope trait. When is the stack space for VariableOps reclaimed when they aren't immediately not surrounded by SPIRV functions? (or are they always are?)

[antiagainst]

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

289

Is it possible to have std.alloca outside of builtin func? So I think we only have Function storage class case here?

296–300

SPIR-V variable ops outside spv.func (that is, with Private storage class) have the lifetime as the SPIR-V module. But is it possible to have std.alloca outside of builtin func? std.alloca's semantics seem to allow that; I'm wondering whether we will see such cases in reality though.

[mravishankar]

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

290

I am not really sure we want to support AllocaOps. The lifetime of the buffers cannot be enforced. So we are potentially breaking that semantics (I dont see a way of deallocating the OpVariable. Maybe some use case info might help

mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp

matchAndRewrite(DeallocOp operation, ArrayRef<Value> operands,
     MemRefType deallocType = operation.memref().getType().cast<MemRefType>();
     if (!isAllocationSupported(deallocType))
       return operation.emitError("unhandled deallocation type");
     rewriter.eraseOp(operation);
     return success();
   }
 };
 
+
+/// Convert Alloca to VariableOp for Function and Private address space.
+class AllocaOpPattern final : public SPIRVOpLowering<AllocaOp> {
+public:
+  using SPIRVOpLowering<AllocaOp>::SPIRVOpLowering;
+
+  LogicalResult
+  matchAndRewrite(AllocaOp operation, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    MemRefType allocType = operation.getType();
+    auto storage = SPIRVTypeConverter::getStorageClassForMemorySpace(
+        allocType.getMemorySpace());
+    if (!allocType.hasStaticShape() || !storage.hasValue())
+      return operation.emitError("unhandled allocation type");
+    if (storage.getValue() != mlir::spirv::StorageClass::Function &&
+        storage.getValue() != mlir::spirv::StorageClass::Private)
+      return operation.emitError("unhandled allocation type");
+    Location loc = operation.getLoc();
+    // Get the SPIR-V type for the allocation.
+    auto spirvType = typeConverter.convertType(allocType);
+    Value alloca = rewriter.create<spirv::VariableOp>(
+        loc, spirvType,
+        rewriter.getI32IntegerAttr(static_cast<int32_t>(storage.getValue())),
+        nullptr);
+    rewriter.replaceOp(operation, alloca);
+    return success();
+  }
+};
+
 /// Converts unary and binary standard operations to SPIR-V operations.
 template <typename StdOp, typename SPIRVOp>
 class UnaryAndBinaryOpPattern final : public SPIRVOpLowering<StdOp> {
 public:
   using SPIRVOpLowering<StdOp>::SPIRVOpLowering;
 
   LogicalResult
   matchAndRewrite(StdOp operation, ArrayRef<Value> operands,
patterns.insert<
       UnaryAndBinaryOpPattern<SinOp, spirv::GLSLSinOp>,
       UnaryAndBinaryOpPattern<SqrtOp, spirv::GLSLSqrtOp>,
       UnaryAndBinaryOpPattern<SubFOp, spirv::FSubOp>,
       UnaryAndBinaryOpPattern<SubIOp, spirv::ISubOp>,
       UnaryAndBinaryOpPattern<TanhOp, spirv::GLSLTanhOp>,
       UnaryAndBinaryOpPattern<UnsignedDivIOp, spirv::UDivOp>,
       UnaryAndBinaryOpPattern<UnsignedRemIOp, spirv::UModOp>,
       UnaryAndBinaryOpPattern<UnsignedShiftRightOp, spirv::ShiftRightLogicalOp>,
-      AllocOpPattern, DeallocOpPattern,
+      AllocOpPattern, DeallocOpPattern, AllocaOpPattern,
       BitwiseOpPattern<AndOp, spirv::LogicalAndOp, spirv::BitwiseAndOp>,
       BitwiseOpPattern<OrOp, spirv::LogicalOrOp, spirv::BitwiseOrOp>,
       BoolCmpIOpPattern, ConstantCompositeOpPattern, ConstantScalarOpPattern,
       CmpFOpPattern, CmpIOpPattern, IntLoadOpPattern, LoadOpPattern,
       ReturnOpPattern, SelectOpPattern, IntStoreOpPattern, StoreOpPattern,
       ZeroExtendI1Pattern, TypeCastingOpPattern<IndexCastOp, spirv::SConvertOp>,
       TypeCastingOpPattern<SIToFPOp, spirv::ConvertSToFOp>,
       TypeCastingOpPattern<ZeroExtendIOp, spirv::UConvertOp>,
       TypeCastingOpPattern<TruncateIOp, spirv::SConvertOp>,
       TypeCastingOpPattern<FPToSIOp, spirv::ConvertFToSOp>,
       TypeCastingOpPattern<FPExtOp, spirv::FConvertOp>,
       TypeCastingOpPattern<FPTruncOp, spirv::FConvertOp>, XOrOpPattern>(
       context, typeConverter);
 }
 } // namespace mlir

mlir/test/Conversion/StandardToSPIRV/alloc.mlir

 {
   func @alloc_dealloc_mem(%arg0 : memref<4x5xf32>) {
     // expected-error @+2 {{unhandled deallocation type}}
     // expected-error @+1 {{op operand #0 must be memref of any type values}}
     dealloc %arg0 : memref<4x5xf32>
     return
   }
 }
+
+
+// -----
+
+module attributes {
+  spv.target_env = #spv.target_env<
+    #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>,
+    {max_compute_workgroup_invocations = 128 : i32,
+     max_compute_workgroup_size = dense<[128, 128, 64]> : vector<3xi32>}>
+  }
+{
+  func @alloca_function_mem(%arg0 : index, %arg1 : index) {
+    %0 = alloca() : memref<4x5xf32, 6>
+    %1 = load %0[%arg0, %arg1] : memref<4x5xf32, 6>
+    store %1, %0[%arg0, %arg1] : memref<4x5xf32, 6>
+    return
+  }
+}
+//     CHECK: func @alloca_function_mem
+// CHECK-NOT:   alloca
+//     CHECK:   %[[PTR:.+]] = spv.Variable : !spv.ptr<!spv.struct<!spv.array<20 x f32, stride=4> [0]>, Function>
+//     CHECK:   %[[LOADPTR:.+]] = spv.AccessChain %[[PTR]]
+//     CHECK:   %[[VAL:.+]] = spv.Load "Function" %[[LOADPTR]] : f32
+//     CHECK:   %[[STOREPTR:.+]] = spv.AccessChain %[[PTR]]
+//     CHECK:   spv.Store "Function" %[[STOREPTR]], %[[VAL]] : f32
+//     CHECK:   spv.Return