[mlir][nvgpu] Add `nvgpu.rcp` OP #100965

Observer007 · 2024-07-29T02:53:24Z

This PR introduces a new OP for reciprocal calculation for vector types using nvvm.rcp OPs. Currently, it supports only f32 types

github-actions · 2024-07-29T02:53:40Z

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be
notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write
permissions for the repository. In which case you can instead tag reviewers by
name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review
by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate
is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

llvmbot · 2024-07-29T02:54:09Z

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-gpu

@llvm/pr-subscribers-mlir-nvgpu

Author: None (Observer007)

Changes

This new op is a simple wrapper op of nvvm.rcp.approx.ftz.f which allows f32 vector input/output.

Full diff: https://github.com/llvm/llvm-project/pull/100965.diff

3 Files Affected:

(modified) mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td (+13)
(modified) mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp (+37-1)
(modified) mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir (+16)

diff --git a/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td b/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td
index dda8f31e688fe..3501c5af16d8c 100644
--- a/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td
+++ b/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td
@@ -20,6 +20,7 @@
 #ifndef NVGPU
 #define NVGPU
 
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/OpBase.td"
@@ -802,4 +803,16 @@ def NVGPU_WarpgroupMmaInitAccumulatorOp : NVGPU_Op<"warpgroup.mma.init.accumulat
   let hasVerifier = 1;
 }
 
+def NVGPU_RcpApproxOp : NVGPU_Op<"rcp_approx", [
+  Pure, SameOperandsAndResultType
+]> {
+  let summary = "A wrapper of nvvm rcp.approx.ftz.f";
+  let description = [{
+    F32 vector reciprocal calculation using `nvvm.rcp.approx.ftz.f`.
+    The input and output are both F32 vector with same shape.
+  }];
+  let arguments = (ins VectorOf<[F32]>:$in);
+  let results = (outs VectorOf<[F32]>:$out);
+  let assemblyFormat = "$in attr-dict `:` type($out)";
+}
 #endif // NVGPU
diff --git a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
index 11d29754aa760..b7ea4aaeb7d8c 100644
--- a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
+++ b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
@@ -11,6 +11,7 @@
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
 #include "mlir/Conversion/LLVMCommon/Pattern.h"
+#include "mlir/Conversion/LLVMCommon/VectorPattern.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
@@ -1666,6 +1667,41 @@ struct NVGPUTmaPrefetchOpLowering
   }
 };
 
+struct NVGPURcpApproxOpLowering
+    : public ConvertOpToLLVMPattern<nvgpu::RcpApproxOp> {
+  using ConvertOpToLLVMPattern<nvgpu::RcpApproxOp>::ConvertOpToLLVMPattern;
+  LogicalResult
+  matchAndRewrite(nvgpu::RcpApproxOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    ImplicitLocOpBuilder b(op->getLoc(), rewriter);
+    auto i64Ty = b.getI64Type();
+    auto f32Ty = b.getF32Type();
+    VectorType inTy = op.getIn().getType();
+    // apply rcp.approx.ftz.f on each element in vector.
+    auto convert1DVec = [&](Type llvm1DVectorTy, Value inVec) {
+      Value ret1DVec = b.create<LLVM::UndefOp>(llvm1DVectorTy);
+      int numElems = llvm::cast<VectorType>(llvm1DVectorTy).getNumElements();
+      for (int i = 0; i < numElems; i++) {
+        Value idx = b.create<LLVM::ConstantOp>(i64Ty, b.getI64IntegerAttr(i));
+        Value elem = b.create<LLVM::ExtractElementOp>(inVec, idx);
+        Value dst = b.create<NVVM::RcpApproxFtzF32Op>(f32Ty, elem);
+        ret1DVec = b.create<LLVM::InsertElementOp>(ret1DVec, dst, idx);
+      }
+      return ret1DVec;
+    };
+    if (inTy.getRank() == 1) {
+      rewriter.replaceOp(op, convert1DVec(inTy, adaptor.getIn()));
+      return success();
+    }
+    return LLVM::detail::handleMultidimensionalVectors(
+        op.getOperation(), adaptor.getOperands(), *(this->getTypeConverter()),
+        [&](Type llvm1DVectorTy, ValueRange operands) -> Value {
+          OpAdaptor adaptor(operands);
+          return convert1DVec(llvm1DVectorTy, adaptor.getIn());
+        },
+        rewriter);
+  }
+};
 } // namespace
 
 void mlir::populateNVGPUToNVVMConversionPatterns(LLVMTypeConverter &converter,
@@ -1688,5 +1724,5 @@ void mlir::populateNVGPUToNVVMConversionPatterns(LLVMTypeConverter &converter,
       NVGPUWarpgroupMmaInitAccumulatorOpLowering, // nvgpu.warpgroup.mma.init.accumulator
       MmaSyncOptoNVVM, MmaLdMatrixOpToNVVM, NVGPUAsyncCopyLowering,
       NVGPUAsyncCreateGroupLowering, NVGPUAsyncWaitLowering,
-      NVGPUMmaSparseSyncLowering>(converter);
+      NVGPUMmaSparseSyncLowering, NVGPURcpApproxOpLowering>(converter);
 }
diff --git a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir
index 86a552c03a473..95cd7a0892b6c 100644
--- a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir
+++ b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir
@@ -1339,3 +1339,19 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// CHECK-LABEL: @rcp_approx_ftz_f32
+// CHECK-SAME:  %[[IN:.*]]: vector<32x16xf32>
+func.func @rcp_approx_ftz_f32(
+    %in: vector<32x16xf32>) {
+  // CHECK: %[[IN_LLVM:.*]] = builtin.unrealized_conversion_cast %[[IN]] : vector<32x16xf32> to !llvm.array<32 x vector<16xf32>>
+  // CHECK: %[[IN1DVEC:.*]] = llvm.extractvalue %[[IN_LLVM]][0] : !llvm.array<32 x vector<16xf32>>
+  // CHECK: %[[OUT1DVEC:.*]] = llvm.mlir.undef : vector<16xf32>
+  // CHECK: %[[IDX_0:.+]] = llvm.mlir.constant(0 : i64) : i64
+  // CHECK: %[[ELEM_0:.*]] = llvm.extractelement %[[IN1DVEC]][%[[IDX_0]] : i64]
+  // CHECK: %[[ELEM_RCP0:.*]] = nvvm.rcp.approx.ftz.f %[[ELEM_0]] : f32
+  // CHECK: llvm.insertelement %[[ELEM_RCP0]], %[[OUT1DVEC]][%[[IDX_0]] : i64] : vector<16xf32>
+  // CHECK-COUNT-511: nvvm.rcp.approx.ftz.f
+  %out = nvgpu.rcp_approx %in : vector<32x16xf32>
+  return
+}

grypp

In general it looks great. Thanks for adding this support in NVGPU dialect. I left some comments

We typically name the titles with a prefix of [mlir] then the dialect [nvgpu]. We can change the title [mlir][nvgpu] Add nvgpu.rcp OP

grypp · 2024-07-29T05:52:02Z

mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td

@@ -20,6 +20,7 @@
 #ifndef NVGPU
 #define NVGPU

+include "mlir/Interfaces/InferTypeOpInterface.td"


Nit : In this dialect, we do not infer types. It is clearer to see the input and output directly in the IR. Otherwise, one would need to read TableGen or, even more challenging, the C++ implementation.

I want to check the input/output type of rcp are equal with SameOperandsAndResultType. So you suggested we'd better change to use verifier?

My point was about inferring the result type. When someone sees an IR like %out = nvgpu.rcp_approx %in : vector<32x16xf32>, the return type is unclear, making the IR unreadable. To understand the return type, one would need to look at the TableGen definition of the operation or read the C++ implementation.

// Current OP, unclear what is %out %out = nvgpu.rcp_approx %in : vector<32x16xf32> // This is my proposal, it's clear what is %out %out = nvgpu.rcp_approx %in : vector<32x16xf32> -> vector<32x16xf32>

You could use PredOpTrait or SameOperandsAndResultType, or handle it in the verifier. Does SameOperandsAndResultType forces to infer result type?

Ok, got it. Let me update the assembly format.

You could use PredOpTrait or SameOperandsAndResultType, or handle it in the verifier. Does SameOperandsAndResultType forces to infer result type?

Yes, SameOperandsAndResultType relies on InferTypeOpInterface and it forces to create build function with InferReturnTypes.

What about PredOpTrait?

Thanks @joker-eph and @grypp , both of formats are good to me. Since most ops of nvgpu dialect have explicit result types, I personally think it is fine to follow this style.

Consistency is important, so that's reasonable.

That said it's not a matter of majority: we can also see the nvgpu dialect as "not using the common upstream practices" and in need of an upgrade.

Recovered original style. Please help check if you have other questions about this pr, thanks a lot. :-)

%out = nvgpu.rcp_approx %in {rounding=approx, ftz}: vector<32x16xf32>

That said it's not a matter of majority: we can also see the nvgpu dialect as "not using the common upstream practices" and in need of an upgrade.

We can upgrade it, but what's the guideline with inferring result types? It makes sense when input and output types same, like with arith. Otherwise, imho, it turns IR obscure and unreadable.

It makes sense when input and output types same, like with arith.

I thought this is the case we're talking about here? The example provided above was %out = nvgpu.rcp_approx %in : vector<32x16xf32> -> vector<32x16xf32>

For the generality, it is slightly more subtle, for example what about %res = memref.load %ptr[%idx] : memref<10xf32> ?

mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td

mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir

grypp

Thanks, looks better

mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td

github-actions · 2024-07-29T10:47:11Z

✅ With the latest revision this PR passed the C/C++ code formatter.

grypp · 2024-07-29T11:56:04Z

@Observer007 you need to run clang-format

github-actions · 2024-07-30T07:21:10Z

@Observer007 Congratulations on having your first Pull Request (PR) merged into the LLVM Project!

Your changes will be combined with recent changes from other authors, then tested
by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR.

Please check whether problems have been caused by your change specifically, as
the builds can include changes from many authors. It is not uncommon for your
change to be included in a build that fails due to someone else's changes, or
infrastructure issues.

How to do this, and the rest of the post-merge process, is covered in detail here.

If your change does cause a problem, it may be reverted, or you can revert it yourself.
This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again.

If you don't get any reports, no action is required from you. Your changes are working as expected, well done!

add rcp approxe op

0749a59

Observer007 requested a review from grypp as a code owner July 29, 2024 02:53

llvmbot added mlir:gpu mlir mlir:nvgpu labels Jul 29, 2024

fix

54d91e8

grypp reviewed Jul 29, 2024

View reviewed changes

Observer007 changed the title ~~[NvGpu Dialect] add rcp approxe op~~ [mlir][nvgpu] add nvgpu.rcp op Jul 29, 2024

Observer007 added 2 commits July 29, 2024 01:58

fix comment

9d1e603

drop useless

9f1829d

grypp approved these changes Jul 29, 2024

View reviewed changes

mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td Outdated Show resolved Hide resolved

mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td Outdated Show resolved Hide resolved

grypp changed the title ~~[mlir][nvgpu] add nvgpu.rcp op~~ [mlir][nvgpu] Add nvgpu.rcp OP Jul 29, 2024

Observer007 added 2 commits July 29, 2024 06:27

apply comment

0094760

recover infer type trait

4b27b91

grypp merged commit 2b23e6c into llvm:main Jul 30, 2024
7 checks passed

[mlir][nvgpu] Add nvgpu.rcp OP #100965

[mlir][nvgpu] Add nvgpu.rcp OP #100965

Uh oh!

Conversation

Observer007 commented Jul 29, 2024 • edited by grypp Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Jul 29, 2024

Uh oh!

llvmbot commented Jul 29, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

grypp left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

grypp left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Jul 29, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

grypp commented Jul 29, 2024

Uh oh!

Uh oh!

github-actions bot commented Jul 30, 2024

Uh oh!

Uh oh!

[mlir][nvgpu] Add `nvgpu.rcp` OP #100965

[mlir][nvgpu] Add `nvgpu.rcp` OP #100965

Observer007 commented Jul 29, 2024 •

edited by grypp

Loading

llvmbot commented Jul 29, 2024 •

edited

Loading

github-actions bot commented Jul 29, 2024 •

edited

Loading