[mlir][gpu] Unroll ops on vectors which map to intrinsic calls

[csigg]

Summary

Unroll ops that map to intrinsics when lowering to LLVM, because intrinsics don't support vector operands/results.

[csigg]

Nothing.

[dcaballe]

Just passing by :). It looks like this is scalarizing a vector operation which sounds like a generic utility that might be worth moving out of GPU? Perhaps VectorUtils.h is a good place?

I'm also wondering if we already have something like that already since we have some vector unrolling transformations to unroll multi-dim vectors into 1-d vectors. Maybe the can be extended and parametrized to how many dims we want to unroll, including all of them? Otherwise, we may want to call this scalarizeVectorOp or something like that?

Hopefully that helps!

[csigg]

Thanks Diego.

I think the vector unroll code should be the generic one, and these rewrites are merely to cope with intrinsics not being available for vectors.

I like scalarizeVectorOp.

[csigg]

Rename to scalarizeVectorOp.

[csigg]

Switch from lambda to template function: uglier formatting, but no uncommon c++ pattern.

[herhut]

mlir/lib/Conversion/GPUCommon/GPUOpsLowering.h

88

nit: operating on

[csigg]

Fix comment.

[nicolasvasilache]

Nice. thanks @csigg !

mlir/lib/Conversion/GPUCommon/GPUOpsLowering.h

struct GPUReturnOpLowering : public ConvertOpToLLVMPattern<gpu::ReturnOp> {
   LogicalResult
   matchAndRewrite(gpu::ReturnOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<LLVM::ReturnOp>(op, adaptor.getOperands());
     return success();
   }
 };
 
+namespace impl {
+/// Unrolls op if it's operating on vectors.
+LogicalResult scalarizeVectorOp(Operation *op, ValueRange operands,
+                                ConversionPatternRewriter &rewriter,
+                                LLVMTypeConverter &converter);
+} // namespace impl
+
+/// Rewriting that unrolls SourceOp to scalars if it's operating on vectors.
+template <typename SourceOp>
+struct ScalarizeVectorOpLowering : public ConvertOpToLLVMPattern<SourceOp> {
+public:
+  using ConvertOpToLLVMPattern<SourceOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(SourceOp op, typename SourceOp::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    return impl::scalarizeVectorOp(op, adaptor.getOperands(), rewriter,
+                                   *this->getTypeConverter());
+  }
+};
+
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_GPUCOMMON_GPUOPSLOWERING_H_

mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp

 //===- GPUOpsLowering.cpp - GPU FuncOp / ReturnOp lowering ----------------===//
 //
 // 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 "GPUOpsLowering.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/FormatVariadic.h"
 
 using namespace mlir;
 
 LogicalResult
 GPUFuncOpLowering::matchAndRewrite(gpu::GPUFuncOp gpuFuncOp, OpAdaptor adaptor,
                                    ConversionPatternRewriter &rewriter) const {
LogicalResult GPUPrintfOpToLLVMCallLowering::matchAndRewrite(
   printfArgs.reserve(argsRange.size() + 1);
   printfArgs.push_back(stringStart);
   printfArgs.append(argsRange.begin(), argsRange.end());
 
   rewriter.create<LLVM::CallOp>(loc, printfDecl, printfArgs);
   rewriter.eraseOp(gpuPrintfOp);
   return success();
 }
+
+/// Unrolls op if it's operating on vectors.
+LogicalResult impl::scalarizeVectorOp(Operation *op, ValueRange operands,
+                                      ConversionPatternRewriter &rewriter,
+                                      LLVMTypeConverter &converter) {
+  TypeRange operandTypes(operands);
+  if (llvm::none_of(operandTypes,
+                    [](Type type) { return type.isa<VectorType>(); })) {
+    return rewriter.notifyMatchFailure(op, "expected vector operand");
+  }
+  if (op->getNumRegions() != 0 || op->getNumSuccessors() != 0)
+    return rewriter.notifyMatchFailure(op, "expected no region/successor");
+  if (op->getNumResults() != 1)
+    return rewriter.notifyMatchFailure(op, "expected single result");
+  VectorType vectorType = op->getResult(0).getType().dyn_cast<VectorType>();
+  if (!vectorType)
+    return rewriter.notifyMatchFailure(op, "expected vector result");
+
+  Location loc = op->getLoc();
+  Value result = rewriter.create<LLVM::UndefOp>(loc, vectorType);
+  Type indexType = converter.convertType(rewriter.getIndexType());
+  StringAttr name = op->getName().getIdentifier();
+  Type elementType = vectorType.getElementType();
+
+  for (int64_t i = 0; i < vectorType.getNumElements(); ++i) {
+    Value index = rewriter.create<LLVM::ConstantOp>(loc, indexType, i);
+    auto extractElement = [&](Value operand) -> Value {
+      if (!operand.getType().isa<VectorType>())
+        return operand;
+      return rewriter.create<LLVM::ExtractElementOp>(loc, operand, index);
+    };
+    auto scalarOperands =
+        llvm::to_vector(llvm::map_range(operands, extractElement));
+    Operation *scalarOp =
+        rewriter.create(loc, name, scalarOperands, elementType, op->getAttrs());
+    rewriter.create<LLVM::InsertElementOp>(loc, result, scalarOp->getResult(0),
+                                           index);
+  }
+
+  rewriter.replaceOp(op, result);
+  return success();
+}

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

 
 #include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
 #include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
 #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h"
 #include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
 #include "mlir/Conversion/LLVMCommon/LoweringOptions.h"
 #include "mlir/Conversion/LLVMCommon/TypeConverter.h"
 #include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
-#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include "llvm/Support/FormatVariadic.h"
 
 #include "../GPUCommon/GPUOpsLowering.h"
 #include "../GPUCommon/IndexIntrinsicsOpLowering.h"
 #include "../GPUCommon/OpToFuncCallLowering.h"
 
 namespace mlir {
 #define GEN_PASS_DEF_CONVERTGPUOPSTONVVMOPS
 #include "mlir/Conversion/Passes.h.inc"
void mlir::configureGpuToNVVMConversionLegality(ConversionTarget &target) {
   target.addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::Exp2Op, LLVM::FAbsOp,
                       LLVM::FCeilOp, LLVM::FFloorOp, LLVM::LogOp, LLVM::Log10Op,
                       LLVM::Log2Op, LLVM::PowOp, LLVM::SinOp, LLVM::SqrtOp>();
 
   // TODO: Remove once we support replacing non-root ops.
   target.addLegalOp<gpu::YieldOp, gpu::GPUModuleOp, gpu::ModuleEndOp>();
 }
 
+template <typename OpTy>
+static void populateOpPatterns(LLVMTypeConverter &converter,
+                               RewritePatternSet &patterns, StringRef f32Func,
+                               StringRef f64Func) {
+  patterns.add<ScalarizeVectorOpLowering<OpTy>>(converter);
+  patterns.add<OpToFuncCallLowering<OpTy>>(converter, f32Func, f64Func);
+}
+
 void mlir::populateGpuToNVVMConversionPatterns(LLVMTypeConverter &converter,
                                                RewritePatternSet &patterns) {
   populateWithGenerated(patterns);
   patterns
       .add<GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, NVVM::ThreadIdXOp,
                                        NVVM::ThreadIdYOp, NVVM::ThreadIdZOp>,
            GPUIndexIntrinsicOpLowering<gpu::BlockDimOp, NVVM::BlockDimXOp,
                                        NVVM::BlockDimYOp, NVVM::BlockDimZOp>,
            GPUIndexIntrinsicOpLowering<gpu::BlockIdOp, NVVM::BlockIdXOp,
                                        NVVM::BlockIdYOp, NVVM::BlockIdZOp>,
            GPUIndexIntrinsicOpLowering<gpu::GridDimOp, NVVM::GridDimXOp,
                                        NVVM::GridDimYOp, NVVM::GridDimZOp>,
            GPULaneIdOpToNVVM, GPUShuffleOpLowering, GPUReturnOpLowering>(
           converter);
 
   // Explicitly drop memory space when lowering private memory
   // attributions since NVVM models it as `alloca`s in the default
   // memory space and does not support `alloca`s with addrspace(5).
   patterns.add<GPUFuncOpLowering>(
       converter, /*allocaAddrSpace=*/0,
       StringAttr::get(&converter.getContext(),
                       NVVM::NVVMDialect::getKernelFuncAttrName()));
 
-  patterns.add<OpToFuncCallLowering<math::AbsFOp>>(converter, "__nv_fabsf",
-                                                   "__nv_fabs");
-  patterns.add<OpToFuncCallLowering<math::AtanOp>>(converter, "__nv_atanf",
-                                                   "__nv_atan");
-  patterns.add<OpToFuncCallLowering<math::Atan2Op>>(converter, "__nv_atan2f",
-                                                    "__nv_atan2");
-  patterns.add<OpToFuncCallLowering<math::CeilOp>>(converter, "__nv_ceilf",
-                                                   "__nv_ceil");
-  patterns.add<OpToFuncCallLowering<math::CosOp>>(converter, "__nv_cosf",
-                                                  "__nv_cos");
-  patterns.add<OpToFuncCallLowering<math::ExpOp>>(converter, "__nv_expf",
-                                                  "__nv_exp");
-  patterns.add<OpToFuncCallLowering<math::Exp2Op>>(converter, "__nv_exp2f",
-                                                   "__nv_exp2");
-  patterns.add<OpToFuncCallLowering<math::ExpM1Op>>(converter, "__nv_expm1f",
-                                                    "__nv_expm1");
-  patterns.add<OpToFuncCallLowering<math::FloorOp>>(converter, "__nv_floorf",
-                                                    "__nv_floor");
-  patterns.add<OpToFuncCallLowering<math::LogOp>>(converter, "__nv_logf",
-                                                  "__nv_log");
-  patterns.add<OpToFuncCallLowering<math::Log1pOp>>(converter, "__nv_log1pf",
-                                                    "__nv_log1p");
-  patterns.add<OpToFuncCallLowering<math::Log10Op>>(converter, "__nv_log10f",
-                                                    "__nv_log10");
-  patterns.add<OpToFuncCallLowering<math::Log2Op>>(converter, "__nv_log2f",
-                                                   "__nv_log2");
-  patterns.add<OpToFuncCallLowering<math::PowFOp>>(converter, "__nv_powf",
-                                                   "__nv_pow");
-  patterns.add<OpToFuncCallLowering<math::RsqrtOp>>(converter, "__nv_rsqrtf",
-                                                    "__nv_rsqrt");
-  patterns.add<OpToFuncCallLowering<math::SinOp>>(converter, "__nv_sinf",
-                                                  "__nv_sin");
-  patterns.add<OpToFuncCallLowering<math::SqrtOp>>(converter, "__nv_sqrtf",
-                                                   "__nv_sqrt");
-  patterns.add<OpToFuncCallLowering<math::TanhOp>>(converter, "__nv_tanhf",
-                                                   "__nv_tanh");
+  populateOpPatterns<math::AbsFOp>(converter, patterns, "__nv_fabsf",
+  populateOpPatterns<math::AtanOp>(converter, patterns, "__nv_atanf",
+  populateOpPatterns<math::Atan2Op>(converter, patterns, "__nv_atan2f",
+  populateOpPatterns<math::CeilOp>(converter, patterns, "__nv_ceilf",
+  populateOpPatterns<math::CosOp>(converter, patterns, "__nv_cosf", "__nv_cos");
+  populateOpPatterns<math::ExpOp>(converter, patterns, "__nv_expf", "__nv_exp");
+  populateOpPatterns<math::Exp2Op>(converter, patterns, "__nv_exp2f",
+  populateOpPatterns<math::ExpM1Op>(converter, patterns, "__nv_expm1f",
+  populateOpPatterns<math::FloorOp>(converter, patterns, "__nv_floorf",
+  populateOpPatterns<math::LogOp>(converter, patterns, "__nv_logf", "__nv_log");
+  populateOpPatterns<math::Log1pOp>(converter, patterns, "__nv_log1pf",
+  populateOpPatterns<math::Log10Op>(converter, patterns, "__nv_log10f",
+  populateOpPatterns<math::Log2Op>(converter, patterns, "__nv_log2f",
+  populateOpPatterns<math::PowFOp>(converter, patterns, "__nv_powf",
+  populateOpPatterns<math::RsqrtOp>(converter, patterns, "__nv_rsqrtf",
+  populateOpPatterns<math::SinOp>(converter, patterns, "__nv_sinf", "__nv_sin");
+  populateOpPatterns<math::SqrtOp>(converter, patterns, "__nv_sqrtf",
+  populateOpPatterns<math::TanhOp>(converter, patterns, "__nv_tanhf",
 }
 
 std::unique_ptr<OperationPass<gpu::GPUModuleOp>>
 mlir::createLowerGpuOpsToNVVMOpsPass(unsigned indexBitwidth) {
   return std::make_unique<LowerGpuOpsToNVVMOpsPass>(indexBitwidth);
 }

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

void mlir::configureGpuToROCDLConversionLegality(ConversionTarget &target) {
   target.addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::Exp2Op, LLVM::FAbsOp,
                       LLVM::FCeilOp, LLVM::FFloorOp, LLVM::LogOp, LLVM::Log10Op,
                       LLVM::Log2Op, LLVM::PowOp, LLVM::SinOp, LLVM::SqrtOp>();
 
   // TODO: Remove once we support replacing non-root ops.
   target.addLegalOp<gpu::YieldOp, gpu::GPUModuleOp, gpu::ModuleEndOp>();
 }
 
+template <typename OpTy>
+static void populateOpPatterns(LLVMTypeConverter &converter,
+                               RewritePatternSet &patterns, StringRef f32Func,
+                               StringRef f64Func) {
+  patterns.add<ScalarizeVectorOpLowering<OpTy>>(converter);
+  patterns.add<OpToFuncCallLowering<OpTy>>(converter, f32Func, f64Func);
+}
+
 void mlir::populateGpuToROCDLConversionPatterns(
     LLVMTypeConverter &converter, RewritePatternSet &patterns,
     mlir::gpu::amd::Runtime runtime) {
   using mlir::gpu::amd::Runtime;
 
   populateWithGenerated(patterns);
   patterns
       .add<GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, ROCDL::ThreadIdXOp,
patterns.add<GPUFuncOpLowering>(
                       ROCDL::ROCDLDialect::getKernelFuncAttrName()));
   if (Runtime::HIP == runtime) {
     patterns.add<GPUPrintfOpToHIPLowering>(converter);
   } else if (Runtime::OpenCL == runtime) {
     // Use address space = 4 to match the OpenCL definition of printf()
     patterns.add<GPUPrintfOpToLLVMCallLowering>(converter, /*addressSpace=*/4);
   }
 
-  patterns.add<OpToFuncCallLowering<math::AbsFOp>>(converter, "__ocml_fabs_f32",
-                                                   "__ocml_fabs_f64");
-  patterns.add<OpToFuncCallLowering<math::AtanOp>>(converter, "__ocml_atan_f32",
-                                                   "__ocml_atan_f64");
-  patterns.add<OpToFuncCallLowering<math::Atan2Op>>(
-      converter, "__ocml_atan2_f32", "__ocml_atan2_f64");
-  patterns.add<OpToFuncCallLowering<math::CeilOp>>(converter, "__ocml_ceil_f32",
-                                                   "__ocml_ceil_f64");
-  patterns.add<OpToFuncCallLowering<math::CosOp>>(converter, "__ocml_cos_f32",
-                                                  "__ocml_cos_f64");
-  patterns.add<OpToFuncCallLowering<math::ExpOp>>(converter, "__ocml_exp_f32",
-                                                  "__ocml_exp_f64");
-  patterns.add<OpToFuncCallLowering<math::Exp2Op>>(converter, "__ocml_exp2_f32",
-                                                   "__ocml_exp2_f64");
-  patterns.add<OpToFuncCallLowering<math::ExpM1Op>>(
-      converter, "__ocml_expm1_f32", "__ocml_expm1_f64");
-  patterns.add<OpToFuncCallLowering<math::FloorOp>>(
-      converter, "__ocml_floor_f32", "__ocml_floor_f64");
-  patterns.add<OpToFuncCallLowering<math::LogOp>>(converter, "__ocml_log_f32",
-                                                  "__ocml_log_f64");
-  patterns.add<OpToFuncCallLowering<math::Log10Op>>(
-      converter, "__ocml_log10_f32", "__ocml_log10_f64");
-  patterns.add<OpToFuncCallLowering<math::Log1pOp>>(
-      converter, "__ocml_log1p_f32", "__ocml_log1p_f64");
-  patterns.add<OpToFuncCallLowering<math::Log2Op>>(converter, "__ocml_log2_f32",
-                                                   "__ocml_log2_f64");
-  patterns.add<OpToFuncCallLowering<math::PowFOp>>(converter, "__ocml_pow_f32",
-                                                   "__ocml_pow_f64");
-  patterns.add<OpToFuncCallLowering<math::RsqrtOp>>(
-      converter, "__ocml_rsqrt_f32", "__ocml_rsqrt_f64");
-  patterns.add<OpToFuncCallLowering<math::SinOp>>(converter, "__ocml_sin_f32",
-                                                  "__ocml_sin_f64");
-  patterns.add<OpToFuncCallLowering<math::SqrtOp>>(converter, "__ocml_sqrt_f32",
-                                                   "__ocml_sqrt_f64");
-  patterns.add<OpToFuncCallLowering<math::TanhOp>>(converter, "__ocml_tanh_f32",
-                                                   "__ocml_tanh_f64");
+  populateOpPatterns<math::AbsFOp>(converter, patterns, "__ocml_fabs_f32",
+  populateOpPatterns<math::AtanOp>(converter, patterns, "__ocml_atan_f32",
+  populateOpPatterns<math::Atan2Op>(converter, patterns, "__ocml_atan2_f32",
+                                    "__ocml_atan2_f64");
+  populateOpPatterns<math::CeilOp>(converter, patterns, "__ocml_ceil_f32",
+  populateOpPatterns<math::CosOp>(converter, patterns, "__ocml_cos_f32",
+  populateOpPatterns<math::ExpOp>(converter, patterns, "__ocml_exp_f32",
+  populateOpPatterns<math::Exp2Op>(converter, patterns, "__ocml_exp2_f32",
+  populateOpPatterns<math::ExpM1Op>(converter, patterns, "__ocml_expm1_f32",
+                                    "__ocml_expm1_f64");
+  populateOpPatterns<math::FloorOp>(converter, patterns, "__ocml_floor_f32",
+                                    "__ocml_floor_f64");
+  populateOpPatterns<math::LogOp>(converter, patterns, "__ocml_log_f32",
+  populateOpPatterns<math::Log10Op>(converter, patterns, "__ocml_log10_f32",
+                                    "__ocml_log10_f64");
+  populateOpPatterns<math::Log1pOp>(converter, patterns, "__ocml_log1p_f32",
+                                    "__ocml_log1p_f64");
+  populateOpPatterns<math::Log2Op>(converter, patterns, "__ocml_log2_f32",
+  populateOpPatterns<math::PowFOp>(converter, patterns, "__ocml_pow_f32",
+  populateOpPatterns<math::RsqrtOp>(converter, patterns, "__ocml_rsqrt_f32",
+                                    "__ocml_rsqrt_f64");
+  populateOpPatterns<math::SinOp>(converter, patterns, "__ocml_sin_f32",
+  populateOpPatterns<math::SqrtOp>(converter, patterns, "__ocml_sqrt_f32",
+  populateOpPatterns<math::TanhOp>(converter, patterns, "__ocml_tanh_f32",
 }
 
 std::unique_ptr<OperationPass<gpu::GPUModuleOp>>
 mlir::createLowerGpuOpsToROCDLOpsPass(const std::string &chipset,
                                       unsigned indexBitwidth,
                                       bool useBarePtrCallConv,
                                       gpu::amd::Runtime runtime) {
   return std::make_unique<LowerGpuOpsToROCDLOpsPass>(
       chipset, indexBitwidth, useBarePtrCallConv, runtime);
 }

mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir

func.func @gpu_pow(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
     // CHECK: llvm.call @__nv_pow(%{{.*}}, %{{.*}}) : (f64, f64) -> f64
     func.return %result32, %result64 : f32, f64
   }
 }
 
 // -----
 
 gpu.module @test_module {
+  // CHECK-LABEL: func @gpu_unroll
+  func.func @gpu_unroll(%arg0 : vector<4xf32>) -> vector<4xf32> {
+    %result = math.exp %arg0 : vector<4xf32>
+    // CHECK: llvm.call @__nv_expf(%{{.*}}) : (f32) -> f32
+    // CHECK: llvm.call @__nv_expf(%{{.*}}) : (f32) -> f32
+    // CHECK: llvm.call @__nv_expf(%{{.*}}) : (f32) -> f32
+    // CHECK: llvm.call @__nv_expf(%{{.*}}) : (f32) -> f32
+    func.return %result : vector<4xf32>
+  }
+}
+
+// -----
+
+gpu.module @test_module {
   // CHECK-LABEL: @kernel_func
   // CHECK: attributes
   // CHECK: gpu.kernel
   // CHECK: nvvm.kernel
   gpu.func @kernel_func() kernel {
     gpu.return
   }
 }

mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir

func.func @gpu_pow(%arg_f32 : f32, %arg_f64 : f64) -> (f32, f64) {
     // CHECK: llvm.call @__ocml_pow_f64(%{{.*}}, %{{.*}}) : (f64, f64) -> f64
     func.return %result32, %result64 : f32, f64
   }
 }
 
 // -----
 
 gpu.module @test_module {
+  // CHECK-LABEL: func @gpu_unroll
+  func.func @gpu_unroll(%arg0 : vector<4xf32>) -> vector<4xf32> {
+    %result = math.exp %arg0 : vector<4xf32>
+    // CHECK: llvm.call @__ocml_exp_f32(%{{.*}}) : (f32) -> f32
+    // CHECK: llvm.call @__ocml_exp_f32(%{{.*}}) : (f32) -> f32
+    // CHECK: llvm.call @__ocml_exp_f32(%{{.*}}) : (f32) -> f32
+    // CHECK: llvm.call @__ocml_exp_f32(%{{.*}}) : (f32) -> f32
+    func.return %result : vector<4xf32>
+  }
+}
+
+// -----
+
+gpu.module @test_module {
   // CHECK-LABEL: @kernel_func
   // CHECK: attributes
   // CHECK: gpu.kernel
   // CHECK: rocdl.kernel
   gpu.func @kernel_func() kernel {
     gpu.return
   }
 }